assn5_vae.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**This is the replication of vae from nn4nlp - 15 vae**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we attempted to train the model from the original dynet implementation, training resulted in an error due to infinite gradients on the 30th epoch.\n",
    "Secondly, the training was done word by word instead of batches being processed simultaneously. \n",
    "Thirdly, vae was trained with a parallel corpus, input data was japanese and the output was english.\n",
    "Lastly, there was no kl divergence annealing or any other method to prevent posterior collapse.\n",
    "\n",
    "In our implementation, we trained for 300 epochs on English to English. To prevent posterior collapse we additionally implemented KL divergence annealing\n",
    "by a linearly incresing coefficient for KL loss. And to speed up the training we trained with minibatches.\n",
    "\n",
    "As a result, in the 2nd to last cell, the pretrained model can be loaded. In the last cell -which works on CPU- a random z from standard gaussian is sampled.\n",
    "Which is used to sample an English sentece form the probability distribution our model learned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Pkg\n",
    "for p in (\"Knet\", \"Random\", \"Statistics\", \"AutoGrad\", \"IterTools\")\n",
    "    haskey(Pkg.installed(),p) || Pkg.add(p)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Comment out last line if using CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Knet, Random, Statistics, Base.Iterators, Test, IterTools\n",
    "using AutoGrad: @gcheck\n",
    "Knet.atype() = KnetArray{Float32}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cells are from previous projects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vocab"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "struct Vocab\n",
    "    w2i::Dict{String,Int}\n",
    "    i2w::Vector{String}\n",
    "    unk::Int\n",
    "    eos::Int\n",
    "    tokenizer\n",
    "end\n",
    "\n",
    "function Vocab(file::String; tokenizer=split, vocabsize=Inf, mincount=1, unk=\"<unk>\", eos=\"<s>\")\n",
    "    M = 100000\n",
    "    wdict = Dict()\n",
    "    wcount = Dict()\n",
    "    w2i(x) = get!(wdict, x, 1+length(wdict))\n",
    "    w2c(key) = haskey(wcount, key) ? wcount[key] = wcount[key] + 1 : get!(wcount, key, 1)\n",
    "    wcount[unk] = M; wcount[eos] = M\n",
    "    i2w = []; \n",
    "\n",
    "    \n",
    "    for line in eachline(file)\n",
    "        words = tokenizer(line)\n",
    "        w2c.(words)\n",
    "    end\n",
    "    \n",
    "    sortedcount = sort(collect(wcount), by=x->x[2])\n",
    "    words = sortedcount[findfirst(x-> x[2]>=mincount, sortedcount):length(sortedcount)]\n",
    "    \n",
    "    #vocabsize excludes unk & eos\n",
    "    if(length(words) > vocabsize)\n",
    "        words = words[length(words) - vocabsize + 1 : length(words)]\n",
    "    end\n",
    "\n",
    "    map(x-> w2i(x[1]) , words)\n",
    "    map(x-> push!(i2w, x[1]), words)\n",
    "    \n",
    "    Vocab(wdict, i2w, wdict[unk], wdict[eos], tokenizer)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct TextReader\n",
    "    file::String\n",
    "    vocab::Vocab\n",
    "end\n",
    "\n",
    "function Base.iterate(r::TextReader, s=nothing)\n",
    "    w2i(x) = get(r.vocab.w2i, x, r.vocab.unk)\n",
    "    if (s === nothing) \n",
    "        s = open(r.file, \"r\")\n",
    "    end\n",
    "\n",
    "    if eof(s) \n",
    "        close(s)\n",
    "        return nothing\n",
    "    \n",
    "    else\n",
    "        tmp = readline(s)\n",
    "        line = r.vocab.tokenizer(tmp)\n",
    "        words = w2i.(line) \n",
    "        return words, s\n",
    "    end    \n",
    "end\n",
    "\n",
    "Base.IteratorSize(::Type{TextReader}) = Base.SizeUnknown()\n",
    "Base.IteratorEltype(::Type{TextReader}) = Base.HasEltype()\n",
    "Base.eltype(::Type{TextReader}) = Vector{Int}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "mask! (generic function with 1 method)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "struct Embed; w; end\n",
    "\n",
    "function Embed(vocabsize::Int, embedsize::Int)\n",
    "    Embed(param(embedsize, vocabsize))\n",
    "end\n",
    "\n",
    "function (l::Embed)(x)\n",
    "    embedsz, vocabsz = size(l.w)\n",
    "    tmparr = [embedsz]\n",
    "    for dim in size(x)\n",
    "        push!(tmparr, dim)\n",
    "    end\n",
    "    reshape(l.w[:,collect(flatten(x))], tuple(tmparr...))\n",
    "end\n",
    "\n",
    "function mask!(a,pad)\n",
    "    x,y = size(a)\n",
    "    \n",
    "    for i = 1:x\n",
    "        tmp_mem = []\n",
    "        isfirst = true\n",
    "        for j = 1:y\n",
    "            if a[i, j] == pad\n",
    "                \n",
    "                if isfirst\n",
    "                    isfirst = false\n",
    "                else\n",
    "                    push!(tmp_mem, j)\n",
    "                end\n",
    "            else\n",
    "                isfirst = true\n",
    "                tmp_mem = []\n",
    "            end\n",
    "        end\n",
    "        tmp_mem = convert(Array{Int,1}, tmp_mem)\n",
    "        a[i, tmp_mem] .= 0\n",
    "    end\n",
    "    return a\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct MTData\n",
    "    src::TextReader        # reader for source language data\n",
    "    tgt::TextReader        # reader for target language data\n",
    "    batchsize::Int         # desired batch size\n",
    "    maxlength::Int         # skip if source sentence above maxlength\n",
    "    batchmajor::Bool       # batch dims (B,T) if batchmajor=false (default) or (T,B) if true.\n",
    "    bucketwidth::Int       # batch sentences with length within bucketwidth of each other\n",
    "    buckets::Vector        # sentences collected in separate arrays called buckets for each length range\n",
    "    batchmaker::Function   # function that turns a bucket into a batch.\n",
    "end\n",
    "\n",
    "#batchsize 128\n",
    "function MTData(src::TextReader, tgt::TextReader; batchmaker = arraybatch, batchsize = 64, maxlength = typemax(Int),\n",
    "                batchmajor = false, bucketwidth = 10, numbuckets = min(128, maxlength ÷ bucketwidth))\n",
    "    buckets = [ [] for i in 1:numbuckets ] # buckets[i] is an array of sentence pairs with similar length\n",
    "    MTData(src, tgt, batchsize, maxlength, batchmajor, bucketwidth, buckets, batchmaker)\n",
    "end\n",
    "\n",
    "Base.IteratorSize(::Type{MTData}) = Base.SizeUnknown()\n",
    "Base.IteratorEltype(::Type{MTData}) = Base.HasEltype()\n",
    "Base.eltype(::Type{MTData}) = NTuple{2}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "function Base.iterate(d::MTData, state=nothing)\n",
    "    if (state === nothing) \n",
    "        \n",
    "        for i = 1:length(d.buckets)\n",
    "            d.buckets[i] = []\n",
    "        end\n",
    "        src = d.src\n",
    "        tgt = d.tgt\n",
    "        src = Iterators.Stateful(src)\n",
    "        tgt = Iterators.Stateful(tgt)\n",
    "    else\n",
    "        src = state[1]\n",
    "        tgt = state[2]\n",
    "    end\n",
    "    \n",
    "    \n",
    "    if(isempty(src)&&isempty(tgt))\n",
    "        for i = 1:length(d.buckets)\n",
    "            if(length(d.buckets[i]) > 0)\n",
    "                tmp_batch = d.batchmaker(d, d.buckets[i])\n",
    "                 if(d.batchmajor == true)\n",
    "                    tmp_batch = (transpose(tmp_batch[1]), transpose(tmp_batch[2]))\n",
    "                end\n",
    "                d.buckets[i] = []\n",
    "                return (tmp_batch, (src, tgt))\n",
    "            end\n",
    "        end\n",
    "    end    \n",
    "        \n",
    "    while(!isempty(src) && !isempty(tgt))\n",
    "        sentences = (popfirst!(src), popfirst!(tgt))\n",
    "        src_sentence = sentences[1]\n",
    "        tgt_sentence = sentences[2]\n",
    "        src_length = length(src_sentence)\n",
    "        \n",
    "        if(src_length > d.maxlength)\n",
    "            continue\n",
    "        elseif(length(d.buckets)*d.bucketwidth < src_length)\n",
    "            index_in_buckets = length(d.buckets)\n",
    "        else\n",
    "            index_in_buckets = ceil(src_length/d.bucketwidth)\n",
    "        end\n",
    "        \n",
    "        index_in_buckets = convert(Int64, index_in_buckets)\n",
    "        push!(d.buckets[index_in_buckets], (src_sentence, tgt_sentence))\n",
    "        \n",
    "        if(isempty(src) && isempty(tgt))\n",
    "                tmp_batch = d.batchmaker(d, d.buckets[index_in_buckets])\n",
    "                if(d.batchmajor == true)\n",
    "                    tmp_batch = (transpose(tmp_batch[1]), transpose(tmp_batch[2]))\n",
    "                end\n",
    "                d.buckets[index_in_buckets] = []\n",
    "                return (tmp_batch, (src, tgt))\n",
    "        end  \n",
    "        \n",
    "        if(length(d.buckets[index_in_buckets]) == d.batchsize)\n",
    "            tmp_batch = d.batchmaker(d, d.buckets[index_in_buckets])\n",
    "            if(d.batchmajor == true)\n",
    "                tmp_batch = (transpose(tmp_batch[1]), transpose(tmp_batch[2]))\n",
    "            end\n",
    "            d.buckets[index_in_buckets] = []\n",
    "            return (tmp_batch, (src, tgt))\n",
    "        end \n",
    "    end   \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "arraybatch (generic function with 1 method)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function arraybatch(d::MTData, bucket)\n",
    "    # Your code here\n",
    "    x = []\n",
    "    y = []\n",
    "    \n",
    "    padded_x = Array{Int64,1}[]\n",
    "    padded_y = Array{Int64,1}[]\n",
    "    \n",
    "    max_length_x = 0\n",
    "    max_length_y = 0\n",
    "    \n",
    "    for sent_pair in bucket\n",
    "        push!(x, sent_pair[1])\n",
    "        push!(sent_pair[2], d.tgt.vocab.eos)\n",
    "        pushfirst!(sent_pair[2], d.tgt.vocab.eos)\n",
    "        push!(y, sent_pair[2])\n",
    "        \n",
    "        if(length(sent_pair[1]) > max_length_x)\n",
    "            max_length_x = length(sent_pair[1])\n",
    "        end\n",
    "        \n",
    "        if(length(sent_pair[2]) > max_length_y)\n",
    "            max_length_y = length(sent_pair[2])\n",
    "        end\n",
    "    end\n",
    "    for sent_pair in zip(x,y)\n",
    "        x_pad_length = max_length_x - length(sent_pair[1])\n",
    "        y_pad_length = max_length_y - length(sent_pair[2])\n",
    "        x_pad_seq = repeat([d.src.vocab.eos], x_pad_length)\n",
    "        y_pad_seq = repeat([d.tgt.vocab.eos], y_pad_length)\n",
    "        push!(padded_x, append!(x_pad_seq, sent_pair[1]))\n",
    "        push!(padded_y, append!(sent_pair[2], y_pad_seq))\n",
    "    end\n",
    "    \n",
    "    no_of_sentences = length(padded_x)\n",
    "\n",
    "    \n",
    "    padded_x = permutedims(hcat(padded_x...), (2,1))\n",
    "    padded_y = permutedims(hcat(padded_y...), (2,1))\n",
    "    \n",
    "    return (padded_x,padded_y)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Data used is from nn4nlp, in our implementation we used only the English corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TextReader(\"nn4nlp-code-master/data/parallel/test.en\", Vocab(Dict(\"enjoy\" => 3043,\"shouldn\" => 1987,\"chocolate\" => 1630,\"fight\" => 2560,\"helping\" => 1988,\"whose\" => 2231,\"hurried\" => 1631,\"favor\" => 2759,\"borders\" => 1,\"star\" => 1632…), [\"borders\", \"stress\", \"fireworks\", \"methods\", \"parted\", \"shakespeare\", \"customer\", \"musical\", \"regarded\", \"21\"  …  \"you\", \"he\", \"is\", \"a\", \"i\", \"to\", \"the\", \".\", \"<unk>\", \"<s>\"], 3710, 3711, split))"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_src_file = \"nn4nlp-code-master/data/parallel/train.ja\"\n",
    "train_tgt_file = \"nn4nlp-code-master/data/parallel/train.en\"\n",
    "dev_src_file = \"nn4nlp-code-master/data/parallel/dev.ja\"\n",
    "dev_tgt_file = \"nn4nlp-code-master/data/parallel/dev.en\"\n",
    "test_src_file = \"nn4nlp-code-master/data/parallel/test.ja\"\n",
    "test_tgt_file = \"nn4nlp-code-master/data/parallel/test.en\"\n",
    "\n",
    "\n",
    "\n",
    "ja_vocab = Vocab(train_src_file, mincount=2)\n",
    "en_vocab = Vocab(train_tgt_file, mincount=2)\n",
    "ja_train = TextReader(train_src_file, ja_vocab)\n",
    "en_train = TextReader(train_tgt_file, en_vocab)\n",
    "ja_dev = TextReader(dev_src_file, ja_vocab)\n",
    "en_dev = TextReader(dev_tgt_file, en_vocab)\n",
    "ja_test = TextReader(test_src_file, ja_vocab)\n",
    "en_test = TextReader(test_tgt_file, en_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"MTData\""
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dtrn = MTData(en_train, en_train)\n",
    "ddev = MTData(en_dev, en_dev)\n",
    "dtst = MTData(en_test, en_test)\n",
    "summary(dtrn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "reparameterize (generic function with 1 method)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function reparameterize(μ, σ)\n",
    "   μ .+ randn!(similar(μ)) .* σ\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct Linear; w; b; end\n",
    "\n",
    "function Linear(inputsize::Int, outputsize::Int)\n",
    "    w = param(outputsize, inputsize)\n",
    "    b = param0(outputsize)\n",
    "    Linear(w,b)\n",
    "end\n",
    "\n",
    "function (l::Linear)(x)\n",
    "    l.w * x .+ l.b\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MLP"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mutable struct MLP\n",
    "    W\n",
    "    V\n",
    "    b\n",
    "end\n",
    "\n",
    "function(m::MLP)(x)\n",
    "    m.V * tanh.(m.W*x .+ m.b)\n",
    "end\n",
    "\n",
    "function MLP(input::Int, hidden::Int, output::Int)\n",
    "    W = param(hidden, input)\n",
    "    b = param0(hidden)\n",
    "    V = param(output, hidden)\n",
    "    MLP(W,V,b)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3711"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Model parameters, identical to the original dynet impplementation\n",
    "EMBED_SIZE = 64\n",
    "HIDDEN_SIZE = 128\n",
    "Q_HIDDEN_SIZE = 64\n",
    "BATCH_SIZE = 16\n",
    "MAX_SENT_SIZE = 50\n",
    "SRC_VOCAB_SIZE = length(ja_vocab.i2w)\n",
    "TGT_VOCAB_SIZE = length(en_vocab.i2w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "mutable struct s2s_vae\n",
    "    srcembed::Embed     # source language embedding\n",
    "    encoder::RNN        # encoder RNN (can be bidirectional)\n",
    "    tgtembed::Embed     # target language embedding\n",
    "    decoder::RNN        # decoder RNN\n",
    "    projection::Linear  # converts decoder output to vocab scores\n",
    "    mean_mlp::MLP       # MLP for estimating mean\n",
    "    var_mlp::MLP        # MLP for estimating standard deviations\n",
    "    srcvocab::Vocab     # source language vocabulary\n",
    "    tgtvocab::Vocab     # target language vocabulary\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "s2s_vae"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function s2s_vae(hidden::Int,        # hidden size for both the encoder and decoder RNN\n",
    "                q_hidden::Int,       # hidden size for MLP hidden layer which estimates mean and std dev for q\n",
    "                srcembsz::Int,       # embedding size for source language\n",
    "                tgtembsz::Int,       # embedding size for target language\n",
    "                srcvocab::Vocab,     # vocabulary for source language\n",
    "                tgtvocab::Vocab;     # vocabulary for target language\n",
    "                layers=1,            # number of layers\n",
    "                bidirectional=false, # whether encoder RNN is bidirectional\n",
    "                ) \n",
    "    \n",
    "    srcembed = Embed(length(srcvocab.i2w), srcembsz)\n",
    "    tgtembed = Embed(length(tgtvocab.i2w), tgtembsz)\n",
    "    \n",
    "    encoder = RNN(srcembsz, hidden, rnnType = :lstm, h = 0)\n",
    "    decoder = RNN(tgtembsz, hidden, rnnType = :lstm, h = 0)\n",
    "    \n",
    "    mean_mlp = MLP(hidden, q_hidden, hidden)\n",
    "    var_mlp = MLP(hidden, q_hidden, hidden)\n",
    "    \n",
    "    projection = Linear(hidden, length(tgtvocab.i2w))\n",
    "    \n",
    "    s2s_vae(srcembed, encoder, tgtembed, decoder, projection, mean_mlp, var_mlp, srcvocab, tgtvocab)\n",
    "    \n",
    "end "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "epoch_no = 0    #for debugging purposes, tracking epoch\n",
    "softmax_loss_total = 0  #for debugging purposes, tracking softmax loss of current epoch\n",
    "kl_loss_total = 0 #for debugging purposes, tracking kl loss of current epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "function (s::s2s_vae)(src, tgt; average=true) #calculate loss for each batch\n",
    "    \n",
    "    global epoch_no\n",
    "    global softmax_loss_total\n",
    "    global kl_loss_total\n",
    "    \n",
    "    #KL loss coefficient increases linearly in the first 80 epochs\n",
    "    linear_divergence_schedule = 80\n",
    "    λ = (linear_divergence_schedule-max(linear_divergence_schedule-epoch_no,0))/linear_divergence_schedule\n",
    "        \n",
    "    #initialize encoder \n",
    "    s.encoder.c = 0; s.encoder.h = 0\n",
    "    \n",
    "    #get the final hidden state from the LSTM given the source sentence \n",
    "    y_enc = s.encoder(s.srcembed(src))[:,:,end]\n",
    "\n",
    "    #estimate means and log variances from the hidden state represenation of the input sentence\n",
    "    mu = s.mean_mlp(y_enc)\n",
    "    log_var = s.var_mlp(y_enc)\n",
    "    \n",
    "    x_mu, y_mu = size(mu)\n",
    "    \n",
    "    #calculate kl loss\n",
    "    kl_loss = -0.5 * sum(1 .+ (log_var - mu.*mu - exp.(log_var)))\n",
    "    \n",
    "    \n",
    "    #perform reparameterization trick\n",
    "    z = reparameterize(mu, exp.(log_var))\n",
    "    x,y = size(z)\n",
    "    z = reshape(z, (x,y,1))\n",
    "    \n",
    "    #initialize decoder according to sampled z\n",
    "    s.decoder.c = z; s.decoder.h = tanh.(z)\n",
    "    y_dec = s.decoder(s.tgtembed(tgt[:,1:end-1]))\n",
    "   \n",
    "    #predict next words with decoder and reconstruction loss by comparing against gold answers\n",
    "    hy, b ,ty = size(y_dec)\n",
    "    y_dec = reshape(y_dec, (hy, b*ty))\n",
    "    scores = s.projection(y_dec)\n",
    "    y_gold = mask!(tgt[:,2:end], s.tgtvocab.eos)\n",
    "    softmax_loss, instances = nll(scores, y_gold; average = false)\n",
    "    \n",
    "    #calculate loss by adding weighted kl loss and reconstruction loss\n",
    "    softmax_loss_total += softmax_loss\n",
    "    kl_loss_total += kl_loss\n",
    "    \n",
    "    λ*kl_loss+softmax_loss\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "loss (generic function with 1 method)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function loss(model, data; average=true) #calculate loss for the entire corpus\n",
    "    instances = 0\n",
    "    loss = 0\n",
    "    global epoch_no\n",
    "    global softmax_loss_total\n",
    "    global kl_loss_total\n",
    "    #track epoch for debugging \n",
    "    epoch_no += 1\n",
    "   \n",
    "    for batch in data\n",
    "        src, tgt = batch\n",
    "        instances += length(tgt)\n",
    "        loss += model(src,tgt)     \n",
    "    end\n",
    "    println(\"softmax loss for epoch \", epoch_no, \" is: \", softmax_loss_total, \" and kl loss is: \", kl_loss_total)\n",
    "    \n",
    "    softmax_loss_total = 0\n",
    "    kl_loss_total = 0\n",
    "    \n",
    "    loss/instances\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "train! (generic function with 1 method)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function train!(model, trn, dev, tst...) #training loop from previous assignments\n",
    "\n",
    "    bestmodel, bestloss = deepcopy(model), loss(model, dev)\n",
    "    progress!(adam(model, trn), steps=100) do y\n",
    "\n",
    "        losses = [ loss(model, d) for d in (dev,tst...) ]\n",
    "        if losses[1] < bestloss\n",
    "            bestmodel, bestloss = deepcopy(model), losses[1]\n",
    "        end\n",
    "        return (losses...,)\n",
    "    end\n",
    "    return bestmodel\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "to initialize model uncomment first, to load a pretrained model uncomment second line\n",
    "to train with parallel japanese to english corpus change the first en_vocab in the first line to ja_vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "s2s_vae(Embed(P(Array{Float32,2}(64,3711))), LSTM(input=64,hidden=128), Embed(P(Array{Float32,2}(64,3711))), LSTM(input=64,hidden=128), Linear(P(Array{Float32,2}(3711,128)), P(Array{Float32,1}(3711))), MLP(P(Array{Float32,2}(64,128)), P(Array{Float32,2}(128,64)), P(Array{Float32,1}(64))), MLP(P(Array{Float32,2}(64,128)), P(Array{Float32,2}(128,64)), P(Array{Float32,1}(64))), Vocab(Dict(\"enjoy\" => 3043,\"shouldn\" => 1987,\"chocolate\" => 1630,\"fight\" => 2560,\"helping\" => 1988,\"whose\" => 2231,\"hurried\" => 1631,\"favor\" => 2759,\"borders\" => 1,\"star\" => 1632…), [\"borders\", \"stress\", \"fireworks\", \"methods\", \"parted\", \"shakespeare\", \"customer\", \"musical\", \"regarded\", \"21\"  …  \"you\", \"he\", \"is\", \"a\", \"i\", \"to\", \"the\", \".\", \"<unk>\", \"<s>\"], 3710, 3711, split), Vocab(Dict(\"enjoy\" => 3043,\"shouldn\" => 1987,\"chocolate\" => 1630,\"fight\" => 2560,\"helping\" => 1988,\"whose\" => 2231,\"hurried\" => 1631,\"favor\" => 2759,\"borders\" => 1,\"star\" => 1632…), [\"borders\", \"stress\", \"fireworks\", \"methods\", \"parted\", \"shakespeare\", \"customer\", \"musical\", \"regarded\", \"21\"  …  \"you\", \"he\", \"is\", \"a\", \"i\", \"to\", \"the\", \".\", \"<unk>\", \"<s>\"], 3710, 3711, split))"
      ]
     },
     "execution_count": 140,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#model = s2s_vae(HIDDEN_SIZE, Q_HIDDEN_SIZE, EMBED_SIZE, EMBED_SIZE, en_vocab, en_vocab)\n",
    "#model = Knet.load(\"vae.jld2\", \"model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11-element Array{Tuple{T,T} where T,1}:\n",
       " ([3947 3947 … 3934 3945; 3947 3947 … 3924 3945; … ; 3947 3947 … 3943 3945; 3947 3947 … 3930 3945], [3711 3685 … 3711 3711; 3711 3704 … 3711 3711; … ; 3711 3710 … 3711 3711; 3711 3703 … 3711 3711])\n",
       " ([3947 3947 … 3940 3945; 3947 3947 … 3925 3945; … ; 3947 3947 … 3940 3945; 3947 3947 … 3930 3945], [3711 3694 … 3711 3711; 3711 3710 … 3711 3711; … ; 3711 3692 … 3711 3711; 3711 3462 … 3711 3711])\n",
       " ([3947 3913 … 3915 3945; 3947 3947 … 3925 3945; … ; 3851 3927 … 3930 3945; 3929 3944 … 3925 3945], [3711 3690 … 3711 3711; 3711 3703 … 3711 3711; … ; 3711 3690 … 3711 3711; 3711 3706 … 3711 3711])\n",
       " ([3947 3947 … 3923 3945; 3947 3947 … 3943 3945; … ; 3947 3947 … 3940 3945; 3947 3947 … 3934 3945], [3711 3702 … 3711 3711; 3711 3691 … 3711 3711; … ; 3711 3690 … 3711 3711; 3711 3290 … 3711 3711])\n",
       " ([3947 3947 … 3934 3945; 3947 3947 … 3940 3945; … ; 3947 3947 … 3943 3945; 3947 3947 … 3940 3945], [3711 3706 … 3711 3711; 3711 2923 … 3711 3711; … ; 3711 2843 … 3711 3711; 3711 3703 … 3711 3711])\n",
       " ([3947 3947 … 3940 3945; 1364 3943 … 3943 3945; … ; 3931 3944 … 3940 3945; 3947 3931 … 3943 3945], [3711 3703 … 3711 3711; 3711 3702 … 3711 3711; … ; 3711 3703 … 3711 3711; 3711 3703 … 3711 3711])\n",
       " ([3947 3235 … 3924 3945; 3947 3947 … 3934 3945; … ; 3822 3941 … 3925 3945; 1644 3944 … 3943 3945], [3711 3704 … 3711 3711; 3711 3708 … 3711 3711; … ; 3711 2862 … 3711 3711; 3711 3706 … 3711 3711])\n",
       " ([3947 3947 … 3923 3945; 3947 3947 … 3943 3945; … ; 3947 3947 … 3925 3945; 3947 3947 … 3934 3945], [3711 3703 … 3711 3711; 3711 3702 … 3711 3711; … ; 3711 3690 … 3711 3711; 3711 3662 … 3711 3711])\n",
       " ([3947 3947 … 3915 3945; 3947 3947 … 3908 3829; … ; 3947 3947 … 3940 3945; 3947 3947 … 3923 3945], [3711 2364 … 3711 3711; 3711 3656 … 3711 3711; … ; 3711 3706 … 3711 3711; 3711 3703 … 3711 3711])\n",
       " ([3947 3947 … 3925 3945; 3840 2387 … 3934 3945; … ; 3947 3947 … 3940 3945; 3947 3947 … 3943 3945], [3711 3694 … 3711 3711; 3711 3482 … 3709 3711; … ; 3711 2834 … 3711 3711; 3711 3700 … 3711 3711])\n",
       " ([3946 3941 … 3934 3945], [3711 3663 … 3709 3711])                                                                                                                                                  "
      ]
     },
     "execution_count": 138,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "epochs = 300  #no of epochs, other parts of the code are from previous assignments\n",
    "ctrn = collect(dtrn)\n",
    "trnx10 = collect(flatten(shuffle!(ctrn) for i in 1:epochs))\n",
    "trn20 = ctrn[1:20]\n",
    "dev38 = collect(ddev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model = train!(model, trnx10, dev38, trn20)   #uncomment to train\n",
    "#Knet.save(\"vae.jld2\",\"model\",model)           #uncomment to save trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "s2s_vae(Embed(P(Array{Float32,2}(64,3711))), LSTM(input=64,hidden=128), Embed(P(Array{Float32,2}(64,3711))), LSTM(input=64,hidden=128), Linear(P(Array{Float32,2}(3711,128)), P(Array{Float32,1}(3711))), MLP(P(Array{Float32,2}(64,128)), P(Array{Float32,2}(128,64)), P(Array{Float32,1}(64))), MLP(P(Array{Float32,2}(64,128)), P(Array{Float32,2}(128,64)), P(Array{Float32,1}(64))), Vocab(Dict(\"enjoy\" => 3043,\"shouldn\" => 1987,\"chocolate\" => 1630,\"fight\" => 2560,\"helping\" => 1988,\"whose\" => 2231,\"hurried\" => 1631,\"favor\" => 2759,\"borders\" => 1,\"star\" => 1632…), [\"borders\", \"stress\", \"fireworks\", \"methods\", \"parted\", \"shakespeare\", \"customer\", \"musical\", \"regarded\", \"21\"  …  \"you\", \"he\", \"is\", \"a\", \"i\", \"to\", \"the\", \".\", \"<unk>\", \"<s>\"], 3710, 3711, split), Vocab(Dict(\"enjoy\" => 3043,\"shouldn\" => 1987,\"chocolate\" => 1630,\"fight\" => 2560,\"helping\" => 1988,\"whose\" => 2231,\"hurried\" => 1631,\"favor\" => 2759,\"borders\" => 1,\"star\" => 1632…), [\"borders\", \"stress\", \"fireworks\", \"methods\", \"parted\", \"shakespeare\", \"customer\", \"musical\", \"regarded\", \"21\"  …  \"you\", \"he\", \"is\", \"a\", \"i\", \"to\", \"the\", \".\", \"<unk>\", \"<s>\"], 3710, 3711, split))"
      ]
     },
     "execution_count": 156,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Knet.load(\"vae.jld2\", \"model\")  #load model for sampling from learned generative model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As mentioned in the beggining, z is sampled from standard gaussian distribution. \n",
    "We use this z as the initial state for our decoder.\n",
    "We sample a new sentence until eos token is reached by greedy decoding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"MersenneTwister\""
      ]
     },
     "execution_count": 210,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = Random.seed!(2)\n",
    "summary(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 211,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9-element Array{String,1}:\n",
       " \"<s>\"   \n",
       " \"i\"     \n",
       " \"have\"  \n",
       " \"to\"    \n",
       " \"go\"    \n",
       " \"to\"    \n",
       " \"school\"\n",
       " \".\"     \n",
       " \"<s>\"   "
      ]
     },
     "execution_count": 211,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#sample from standard gaussian\n",
    "z = reshape(reparameterize(zeros(Float32, 128),ones(Float32, 128)), (128,1,1))  \n",
    "\n",
    "#inititalize decoder\n",
    "model.decoder.h = tanh.(z)\n",
    "model.decoder.c = z\n",
    "input = [model.tgtvocab.eos]\n",
    "isDone = false\n",
    "translated_sentence = [model.tgtvocab.eos]\n",
    "    \n",
    "#sample next words with greedy decoding until eos is reached\n",
    "while (!isDone)\n",
    "        global input\n",
    "        input = reshape(input, (1,1))\n",
    "        y = model.decoder(model.tgtembed(input))\n",
    "        \n",
    "        scores = model.projection(mat(y))\n",
    "        next_word = argmax(scores)[1]\n",
    "        translated_sentence = push!(translated_sentence, next_word)\n",
    "        input = [next_word]\n",
    "        if(next_word == model.tgtvocab.eos || length(translated_sentence)>50)\n",
    "            isDone = true\n",
    "        end\n",
    "    \n",
    "end\n",
    "#print out the final sentence\n",
    "model.tgtvocab.i2w[translated_sentence]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.2.0",
   "language": "julia",
   "name": "julia-1.2"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}