assn4_attn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attention-based Neural Machine Translation\n",
    "\n",
    "**Reference:** Luong, Thang, Hieu Pham and Christopher D. Manning. \"Effective Approaches to Attention-based Neural Machine Translation.\" In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pp. 1412-1421. 2015.\n",
    "\n",
    "* https://www.aclweb.org/anthology/D15-1166/ (main paper reference)\n",
    "* https://arxiv.org/abs/1508.04025 (alternative paper url)\n",
    "* https://github.com/tensorflow/nmt (main code reference)\n",
    "* https://www.tensorflow.org/beta/tutorials/text/nmt_with_attention (alternative code reference)\n",
    "* https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py:2449,2103 (attention implementation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m registry at `~/.julia/registries/General`\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m git-repo `https://github.com/JuliaRegistries/General.git`\n",
      "\u001b[?25l\u001b[2K\u001b[?25h\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(nothing, nothing)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import Pkg\n",
    "using Pkg\n",
    "Pkg.add(\"Knet\"); Pkg.add(\"CuArrays\"), Pkg.add(\"Random\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "┌ Info: Precompiling CuArrays [3a865a2d-5b23-5a0f-bc46-62713ec82fae]\n",
      "└ @ Base loading.jl:1242\n",
      "ERROR: LoadError: Could not find CUDA driver library\n",
      "Stacktrace:\n",
      " [1] error(::String) at ./error.jl:33\n",
      " [2] top-level scope at /Users/deniz/.julia/packages/CUDAdrv/ADRHQ/src/CUDAdrv.jl:33\n",
      " [3] include at ./boot.jl:328 [inlined]\n",
      " [4] include_relative(::Module, ::String) at ./loading.jl:1094\n",
      " [5] include(::Module, ::String) at ./Base.jl:31\n",
      " [6] top-level scope at none:2\n",
      " [7] eval at ./boot.jl:330 [inlined]\n",
      " [8] eval(::Expr) at ./client.jl:432\n",
      " [9] top-level scope at ./none:3\n",
      "in expression starting at /Users/deniz/.julia/packages/CUDAdrv/ADRHQ/src/CUDAdrv.jl:27\n",
      "ERROR: LoadError: Failed to precompile CUDAdrv [c5f51814-7f29-56b8-a69c-e4d8f6be1fde] to /Users/deniz/.julia/compiled/v1.2/CUDAdrv/HMhfu.ji.\n",
      "Stacktrace:\n",
      " [1] error(::String) at ./error.jl:33\n",
      " [2] compilecache(::Base.PkgId, ::String) at ./loading.jl:1253\n",
      " [3] _require(::Base.PkgId) at ./loading.jl:1013\n",
      " [4] require(::Base.PkgId) at ./loading.jl:911\n",
      " [5] require(::Module, ::Symbol) at ./loading.jl:906\n",
      " [6] include at ./boot.jl:328 [inlined]\n",
      " [7] include_relative(::Module, ::String) at ./loading.jl:1094\n",
      " [8] include(::Module, ::String) at ./Base.jl:31\n",
      " [9] top-level scope at none:2\n",
      " [10] eval at ./boot.jl:330 [inlined]\n",
      " [11] eval(::Expr) at ./client.jl:432\n",
      " [12] top-level scope at ./none:3\n",
      "in expression starting at /Users/deniz/.julia/packages/CuArrays/kOUu1/src/CuArrays.jl:3\n"
     ]
    },
    {
     "ename": "ErrorException",
     "evalue": "Failed to precompile CuArrays [3a865a2d-5b23-5a0f-bc46-62713ec82fae] to /Users/deniz/.julia/compiled/v1.2/CuArrays/7YFE0.ji.",
     "output_type": "error",
     "traceback": [
      "Failed to precompile CuArrays [3a865a2d-5b23-5a0f-bc46-62713ec82fae] to /Users/deniz/.julia/compiled/v1.2/CuArrays/7YFE0.ji.",
      "",
      "Stacktrace:",
      " [1] error(::String) at ./error.jl:33",
      " [2] compilecache(::Base.PkgId, ::String) at ./loading.jl:1253",
      " [3] _require(::Base.PkgId) at ./loading.jl:1013",
      " [4] require(::Base.PkgId) at ./loading.jl:911",
      " [5] require(::Module, ::Symbol) at ./loading.jl:906",
      " [6] top-level scope at In[2]:2"
     ]
    }
   ],
   "source": [
    "using Knet, Test, Base.Iterators, Printf, LinearAlgebra, Random, IterTools\n",
    "using CuArrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code and data from previous projects\n",
    "\n",
    "Please copy or include the following types and related functions from previous projects:\n",
    "`Vocab`, `TextReader`, `MTData`, `Embed`, `Linear`, `mask!`, `loss`, `int2str`,\n",
    "`bleu`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "bleu (generic function with 1 method)"
      ]
     },
     "execution_count": 90,
     "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\n",
    "\n",
    "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}\n",
    "\n",
    "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",
    "\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",
    "\n",
    "end\n",
    "\n",
    "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\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\n",
    "\n",
    "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}\n",
    "\n",
    "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\n",
    "\n",
    "\n",
    "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\n",
    "\n",
    "function loss(model, data; average=true)\n",
    "    instances = 0\n",
    "    cumulative_loss = 0\n",
    "    for batch in data\n",
    "        x, y = batch\n",
    "        batch_loss, batch_instances = model(x,y; average=false)\n",
    "        cumulative_loss += batch_loss\n",
    "        instances += batch_instances\n",
    "    end\n",
    "    if (average)\n",
    "        cumulative_loss / instances\n",
    "    else\n",
    "        cumulative_loss, instances\n",
    "    end\n",
    "end\n",
    "\n",
    "function int2str(y,vocab)\n",
    "    y = vec(y)\n",
    "    ysos = findnext(w->!isequal(w,vocab.eos), y, 1)\n",
    "    ysos == nothing && return \"\"\n",
    "    yeos = something(findnext(isequal(vocab.eos), y, ysos), 1+length(y))\n",
    "    join(vocab.i2w[y[ysos:yeos-1]], \" \")\n",
    "end\n",
    "\n",
    "\n",
    "function bleu(s2s,d::MTData)\n",
    "    d = MTData(d.src,d.tgt,batchsize=1)\n",
    "    reffile = d.tgt.file\n",
    "    hypfile,hyp = mktemp()\n",
    "    for (x,y) in progress(collect(d))\n",
    "        g = s2s(x)\n",
    "        for i in 1:size(y,1)\n",
    "            println(hyp, int2str(g[i,:], d.tgt.vocab))\n",
    "        end\n",
    "    end\n",
    "    close(hyp)\n",
    "    isfile(\"multi-bleu.perl\") || download(\"https://github.com/moses-smt/mosesdecoder/raw/master/scripts/generic/multi-bleu.perl\", \"multi-bleu.perl\")\n",
    "    run(pipeline(`cat $hypfile`,`perl multi-bleu.perl $reffile`))\n",
    "    return hypfile\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## S2S: Sequence to sequence model with attention\n",
    "\n",
    "In this project we will define, train and evaluate a sequence to sequence encoder-decoder\n",
    "model with attention for Turkish-English machine translation. The model has two extra\n",
    "fields compared to `S2S_v1`: the `memory` layer computes keys and values from the encoder,\n",
    "the `attention` layer computes the attention vector for the decoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct Memory; w; end\n",
    "\n",
    "struct Attention; wquery; wattn; scale; end\n",
    "\n",
    "struct S2S\n",
    "    srcembed::Embed       # encinput(B,Tx) -> srcembed(Ex,B,Tx)\n",
    "    encoder::RNN          # srcembed(Ex,B,Tx) -> enccell(Dx*H,B,Tx)\n",
    "    memory::Memory        # enccell(Dx*H,B,Tx) -> keys(H,Tx,B), vals(Dx*H,Tx,B)\n",
    "    tgtembed::Embed       # decinput(B,Ty) -> tgtembed(Ey,B,Ty)\n",
    "    decoder::RNN          # tgtembed(Ey,B,Ty) . attnvec(H,B,Ty)[t-1] = (Ey+H,B,Ty) -> deccell(H,B,Ty)\n",
    "    attention::Attention  # deccell(H,B,Ty), keys(H,Tx,B), vals(Dx*H,Tx,B) -> attnvec(H,B,Ty)\n",
    "    projection::Linear    # attnvec(H,B,Ty) -> proj(Vy,B,Ty)\n",
    "    dropout::Real         # dropout probability\n",
    "    srcvocab::Vocab       # source language vocabulary\n",
    "    tgtvocab::Vocab       # target language vocabulary\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load pretrained model and data\n",
    "\n",
    "We will load a pretrained model (16.20 bleu) for code testing.  The data should be loaded\n",
    "with the vocabulary from the pretrained model for word id consistency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "if !isdefined(Main, :pretrained) || pretrained === nothing\n",
    "    @info \"Loading reference model\"\n",
    "    isfile(\"s2smodel.jld2\") || download(\"http://people.csail.mit.edu/deniz/comp542/s2smodel.jld2\",\"s2smodel.jld2\")\n",
    "    pretrained = Knet.load(\"s2smodel.jld2\",\"model\")\n",
    "end\n",
    "datadir = \"datasets/tr_to_en\"\n",
    "if !isdir(datadir)\n",
    "    @info \"Downloading data\"\n",
    "    download(\"http://www.phontron.com/data/qi18naacl-dataset.tar.gz\", \"qi18naacl-dataset.tar.gz\")\n",
    "    run(`tar xzf qi18naacl-dataset.tar.gz`)\n",
    "end\n",
    "if !isdefined(Main, :tr_vocab)\n",
    "    BATCHSIZE, MAXLENGTH = 64, 50\n",
    "    @info \"Reading data\"\n",
    "    tr_vocab = pretrained.srcvocab # Vocab(\"$datadir/tr.train\", mincount=5)\n",
    "    en_vocab = pretrained.tgtvocab # Vocab(\"$datadir/en.train\", mincount=5)\n",
    "    tr_train = TextReader(\"$datadir/tr.train\", tr_vocab)\n",
    "    en_train = TextReader(\"$datadir/en.train\", en_vocab)\n",
    "    tr_dev = TextReader(\"$datadir/tr.dev\", tr_vocab)\n",
    "    en_dev = TextReader(\"$datadir/en.dev\", en_vocab)\n",
    "    tr_test = TextReader(\"$datadir/tr.test\", tr_vocab)\n",
    "    en_test = TextReader(\"$datadir/en.test\", en_vocab)\n",
    "    dtrn = MTData(tr_train, en_train, batchsize=BATCHSIZE, maxlength=MAXLENGTH)\n",
    "    ddev = MTData(tr_dev, en_dev, batchsize=BATCHSIZE)\n",
    "    dtst = MTData(tr_test, en_test, batchsize=BATCHSIZE)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1. Model constructor\n",
    "\n",
    "The `S2S` constructor takes the following arguments:\n",
    "* `hidden`: size of the hidden vectors for both the encoder and the decoder\n",
    "* `srcembsz`, `tgtembsz`: size of the source/target language embedding vectors\n",
    "* `srcvocab`, `tgtvocab`: the source/target language vocabulary\n",
    "* `layers=1`: number of layers\n",
    "* `bidirectional=false`: whether the encoder is bidirectional\n",
    "* `dropout=0`: dropout probability\n",
    "\n",
    "Hints:\n",
    "* You can find the vocabulary size with `length(vocab.i2w)`.\n",
    "* If the encoder is bidirectional `layers` must be even and the encoder should have `layers÷2` layers.\n",
    "* The decoder will use \"input feeding\", i.e. it will concatenate its previous output to its input. Therefore the input size for the decoder should be `tgtembsz+hidden`.\n",
    "* Only `numLayers`, `dropout`, and `bidirectional` keyword arguments should be used for RNNs, leave everything else default.\n",
    "* The memory parameter `w` is used to convert encoder states to keys. If the encoder is bidirectional initialize it to a `(hidden,2*hidden)` parameter, otherwise set it to the constant 1.\n",
    "* The attention parameter `wquery` is used to transform the query, set it to the constant 1 for this project.\n",
    "* The attention parameter `scale` is used to scale the attention scores before softmax, set it to a parameter of size 1.\n",
    "* The attention parameter `wattn` is used to transform the concatenation of the decoder output and the context vector to the attention vector. It should be a parameter of size `(hidden,2*hidden)` if unidirectional, `(hidden,3*hidden)` if bidirectional."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "S2S"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function S2S(hidden::Int, srcembsz::Int, tgtembsz::Int, srcvocab::Vocab, tgtvocab::Vocab;\n",
    "             layers=1, bidirectional=false, dropout=0)\n",
    "    \n",
    "    srcembed = Embed(length(srcvocab.i2w), srcembsz)\n",
    "    tgtembed = Embed(length(tgtvocab.i2w), tgtembsz)\n",
    "    decoder_layers = layers\n",
    "    memory_w = 1\n",
    "    attn_wq = 1\n",
    "    attn_scale = param(1)\n",
    "    wattn = Linear(hidden, 2*hidden)\n",
    "    if(bidirectional == true)\n",
    "        encoder_layers = layers/2\n",
    "        memory_w = param(hidden, 2*hidden)\n",
    "        wattn = param(hidden, 3*hidden)\n",
    "    end\n",
    "    memory = Memory(memory_w)\n",
    "    attention = Attention(attn_wq, wattn, attn_scale)\n",
    "    \n",
    "    \n",
    "    encoder = RNN(srcembsz, hidden, rnnType = :lstm, bidirectional = bidirectional, dropout = dropout, numLayers = encoder_layers, h = 0)\n",
    "    decoder = RNN(tgtembsz+hidden, hidden, rnnType = :lstm, dropout = dropout, numLayers = layers, h = 0)\n",
    "    projection = Linear(hidden, length(tgtvocab.i2w))\n",
    "    \n",
    "    S2S(srcembed, encoder, memory, tgtembed, decoder, attention, projection, dropout, srcvocab, tgtvocab)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37m\u001b[1mTest Summary:           | \u001b[22m\u001b[39m\u001b[32m\u001b[1mPass  \u001b[22m\u001b[39m\u001b[36m\u001b[1mTotal\u001b[22m\u001b[39m\n",
      "Testing S2S constructor | \u001b[32m  16  \u001b[39m\u001b[36m   16\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Test.DefaultTestSet(\"Testing S2S constructor\", Any[], 16, false)"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@testset \"Testing S2S constructor\" begin\n",
    "    H,Ex,Ey,Vx,Vy,L,Dx,Pdrop = 8,9,10,length(dtrn.src.vocab.i2w),length(dtrn.tgt.vocab.i2w),2,2,0.2\n",
    "    m = S2S(H,Ex,Ey,dtrn.src.vocab,dtrn.tgt.vocab;layers=L,bidirectional=(Dx==2),dropout=Pdrop)\n",
    "    @test size(m.srcembed.w) == (Ex,Vx)\n",
    "    @test size(m.tgtembed.w) == (Ey,Vy)\n",
    "    @test m.encoder.inputSize == Ex\n",
    "    @test m.decoder.inputSize == Ey + H\n",
    "    @test m.encoder.hiddenSize == m.decoder.hiddenSize == H\n",
    "    @test m.encoder.direction == Dx-1\n",
    "    @test m.encoder.numLayers == (Dx == 2 ? L÷2 : L)\n",
    "    @test m.decoder.numLayers == L\n",
    "    @test m.encoder.dropout == m.decoder.dropout == Pdrop\n",
    "    @test size(m.projection.w) == (Vy,H)\n",
    "    @test size(m.memory.w) == (Dx == 2 ? (H,2H) : ())\n",
    "    @test m.attention.wquery == 1\n",
    "    @test size(m.attention.wattn) == (Dx == 2 ? (H,3H) : (H,2H))\n",
    "    @test size(m.attention.scale) == (1,)\n",
    "    @test m.srcvocab === dtrn.src.vocab\n",
    "    @test m.tgtvocab === dtrn.tgt.vocab\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2. Memory\n",
    "\n",
    "The memory layer turns the output of the encoder to a pair of tensors that will be used as\n",
    "keys and values for the attention mechanism. Remember that the encoder RNN output has size\n",
    "`(H*D,B,Tx)` where `H` is the hidden size, `D` is 1 for unidirectional, 2 for\n",
    "bidirectional, `B` is the batchsize, and `Tx` is the sequence length. It will be\n",
    "convenient to store these values in batch major form for the attention mechanism, so\n",
    "*values* in memory will be a permuted copy of the encoder output with size `(H*D,Tx,B)`\n",
    "(see `@doc permutedims`). The *keys* in the memory need to have the same first dimension\n",
    "as the *queries* (i.e. the decoder hidden states). So *values* will be transformed into\n",
    "*keys* of size `(H,B,Tx)` with `keys = m.w * values` where `m::Memory` is the memory\n",
    "layer. Note that you will have to do some reshaping to 2-D and back to 3-D for matrix\n",
    "multiplications. Also note that `m.w` may be a scalar such as `1` e.g. when `D=1` and we\n",
    "want keys and values to be identical."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "function (m::Memory)(x)\n",
    "    vals = permutedims(x, (1,3,2))\n",
    "    keys = mmul(m.w, vals)\n",
    "    return keys, vals\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use the following helper function for scaling and linear transformations of 3-D tensors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "mmul (generic function with 1 method)"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mmul(w,x) = (w == 1 ? x : w == 0 ? 0 : reshape(w * reshape(x,size(x,1),:), (:, size(x)[2:end]...)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37mTesting memory: \u001b[39m\u001b[91m\u001b[1mError During Test\u001b[22m\u001b[39m at \u001b[39m\u001b[1mIn[97]:1\u001b[22m\n",
      "  Got exception outside of a @test\n",
      "  TypeError: in Type{...} expression, expected UnionAll, got typeof(Knet.CuArray)\n",
      "  Stacktrace:\n",
      "   [1] KnetPtrCu(::Int64) at /Users/deniz/.julia/packages/Knet/HRYiN/src/cuarray.jl:90\n",
      "   [2] Knet.KnetPtr(::Int64) at /Users/deniz/.julia/packages/Knet/HRYiN/src/kptr.jl:102\n",
      "   [3] KnetArray{Float32,N} where N(::UndefInitializer, ::Tuple{Int64,Int64,Int64}) at /Users/deniz/.julia/packages/Knet/HRYiN/src/karray.jl:82\n",
      "   [4] KnetArray{Float32,3}(::Array{Float32,3}) at /Users/deniz/.julia/packages/Knet/HRYiN/src/karray.jl:95\n",
      "   [5] KnetArray(::Array{Float32,3}) at /Users/deniz/.julia/packages/Knet/HRYiN/src/karray.jl:93\n",
      "   [6] top-level scope at In[97]:3\n",
      "   [7] top-level scope at /Users/sabae/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.2/Test/src/Test.jl:1113\n",
      "   [8] top-level scope at In[97]:2\n",
      "   [9] eval at ./boot.jl:330 [inlined]\n",
      "   [10] softscope_include_string(::Module, ::String, ::String) at /Users/deniz/.julia/packages/SoftGlobalScope/cSbw5/src/SoftGlobalScope.jl:218\n",
      "   [11] execute_request(::ZMQ.Socket, ::IJulia.Msg) at /Users/deniz/.julia/packages/IJulia/F1GUo/src/execute_request.jl:67\n",
      "   [12] #invokelatest#1 at ./essentials.jl:790 [inlined]\n",
      "   [13] invokelatest at ./essentials.jl:789 [inlined]\n",
      "   [14] eventloop(::ZMQ.Socket) at /Users/deniz/.julia/packages/IJulia/F1GUo/src/eventloop.jl:8\n",
      "   [15] (::getfield(IJulia, Symbol(\"##15#18\")))() at ./task.jl:268\n",
      "  \n",
      "\u001b[37m\u001b[1mTest Summary:  | \u001b[22m\u001b[39m\u001b[91m\u001b[1mError  \u001b[22m\u001b[39m\u001b[36m\u001b[1mTotal\u001b[22m\u001b[39m\n",
      "Testing memory | \u001b[91m    1  \u001b[39m\u001b[36m    1\u001b[39m\n"
     ]
    },
    {
     "ename": "TestSetException",
     "evalue": "Some tests did not pass: 0 passed, 0 failed, 1 errored, 0 broken.",
     "output_type": "error",
     "traceback": [
      "Some tests did not pass: 0 passed, 0 failed, 1 errored, 0 broken.",
      "",
      "Stacktrace:",
      " [1] finish(::Test.DefaultTestSet) at /Users/sabae/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.2/Test/src/Test.jl:877",
      " [2] top-level scope at /Users/sabae/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.2/Test/src/Test.jl:1123",
      " [3] top-level scope at In[97]:2"
     ]
    }
   ],
   "source": [
    "@testset \"Testing memory\" begin\n",
    "    H,D,B,Tx = pretrained.encoder.hiddenSize, pretrained.encoder.direction+1, 4, 5\n",
    "    x = KnetArray(randn(Float32,H*D,B,Tx))\n",
    "    k,v = pretrained.memory(x)\n",
    "    @test v == permutedims(x,(1,3,2))\n",
    "    @test k == mmul(pretrained.memory.w, v)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3. Encoder\n",
    "\n",
    "`encode()` takes a model `s` and a source language minibatch `src`. It passes the input\n",
    "through `s.srcembed` and `s.encoder` layers with the `s.encoder` RNN hidden states\n",
    "initialized to `0` in the beginning, and copied to the `s.decoder` RNN at the end. The\n",
    "steps so far are identical to `S2S_v1` but there is an extra step: The encoder output is\n",
    "passed to the `s.memory` layer which returns a `(keys,values)` pair. `encode()` returns\n",
    "this pair to be used later by the attention mechanism."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "encode (generic function with 1 method)"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function encode(s::S2S, src)\n",
    "    src_embed_tensor = dropout(s.srcembed(src), s.dropout)\n",
    "    s.encoder.h = 0\n",
    "    s.encoder.c = 0\n",
    "    y_enc = s.encoder(src_embed_tensor)\n",
    "    s.decoder.h = s.encoder.h\n",
    "    s.decoder.c = s.encoder.c\n",
    "    \n",
    "    keys, values = s.memory(y_enc)\n",
    "    return keys, values\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37m\u001b[1mTest Summary:   | \u001b[22m\u001b[39m\u001b[32m\u001b[1mPass  \u001b[22m\u001b[39m\u001b[36m\u001b[1mTotal\u001b[22m\u001b[39m\n",
      "Testing encoder | \u001b[32m   7  \u001b[39m\u001b[36m    7\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Test.DefaultTestSet(\"Testing encoder\", Any[], 7, false)"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@testset \"Testing encoder\" begin\n",
    "    src1,tgt1 = first(dtrn)\n",
    "    key1,val1 = encode(pretrained, src1)\n",
    "    H,D,B,Tx = pretrained.encoder.hiddenSize, pretrained.encoder.direction+1, size(src1,1), size(src1,2)\n",
    "    @test size(key1) == (H,Tx,B)\n",
    "    @test size(val1) == (H*D,Tx,B)\n",
    "    @test (pretrained.decoder.h,pretrained.decoder.c) === (pretrained.encoder.h,pretrained.encoder.c)\n",
    "    @test norm(key1) ≈ 1214.4755f0\n",
    "    @test norm(val1) ≈ 191.10411f0\n",
    "    @test norm(pretrained.decoder.h) ≈ 48.536964f0\n",
    "    @test norm(pretrained.decoder.c) ≈ 391.69028f0\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 4. Attention\n",
    "\n",
    "The attention layer takes `cell`: the decoder output, and `mem`: a pair of (keys,vals)\n",
    "from the encoder, and computes and returns the attention vector. First `a.wquery` is used\n",
    "to linearly transform the cell to the query tensor. The query tensor is reshaped and/or\n",
    "permuted as appropriate and multiplied with the keys tensor to compute the attention\n",
    "scores. Please see `@doc bmm` for the batched matrix multiply operation used for this\n",
    "step. The attention scores are scaled using `a.scale` and normalized along the time\n",
    "dimension using `softmax`. After the appropriate reshape and/or permutation, the scores\n",
    "are multiplied with the `vals` tensor (using `bmm` again) to compute the context\n",
    "tensor. After the appropriate reshape and/or permutation the context vector is\n",
    "concatenated with the cell and linearly transformed to the attention vector using\n",
    "`a.wattn`. Please see the paper and code examples for details.\n",
    "\n",
    "Note: the paper mentions a final `tanh` transform, however the final version of the\n",
    "reference code does not use `tanh` and gets better results. Therefore we will skip `tanh`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "function (a::Attention)(cell, mem)\n",
    "    keys, values = mem\n",
    "    query = permutedims(mmul(a.wquery, cell), (3,1,2))\n",
    "    scores = bmm(query, keys)\n",
    "    scores = mmul(a.scale[1], scores)\n",
    "    \n",
    "    scores = softmax(scores, dims = 2)\n",
    "    context = bmm(values, permutedims(scores, (2,1,3)))\n",
    "    mmul(a.wattn, vcat(cell,permutedims(context, (1,3,2))))\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37mTesting attention: \u001b[39m\u001b[91m\u001b[1mError During Test\u001b[22m\u001b[39m at \u001b[39m\u001b[1mIn[101]:1\u001b[22m\n",
      "  Got exception outside of a @test\n",
      "  TypeError: in Type{...} expression, expected UnionAll, got typeof(Knet.CuArray)\n",
      "  Stacktrace:\n",
      "   [1] KnetPtrCu(::Int64) at /Users/deniz/.julia/packages/Knet/HRYiN/src/cuarray.jl:90\n",
      "   [2] Knet.KnetPtr(::Int64) at /Users/deniz/.julia/packages/Knet/HRYiN/src/kptr.jl:102\n",
      "   [3] KnetArray{Float32,N} where N(::UndefInitializer, ::Tuple{Int64,Int64,Int64}) at /Users/deniz/.julia/packages/Knet/HRYiN/src/karray.jl:82\n",
      "   [4] KnetArray{Float32,3}(::Array{Float32,3}) at /Users/deniz/.julia/packages/Knet/HRYiN/src/karray.jl:95\n",
      "   [5] KnetArray(::Array{Float32,3}) at /Users/deniz/.julia/packages/Knet/HRYiN/src/karray.jl:93\n",
      "   [6] top-level scope at In[101]:6\n",
      "   [7] top-level scope at /Users/sabae/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.2/Test/src/Test.jl:1113\n",
      "   [8] top-level scope at In[101]:2\n",
      "   [9] eval at ./boot.jl:330 [inlined]\n",
      "   [10] softscope_include_string(::Module, ::String, ::String) at /Users/deniz/.julia/packages/SoftGlobalScope/cSbw5/src/SoftGlobalScope.jl:218\n",
      "   [11] execute_request(::ZMQ.Socket, ::IJulia.Msg) at /Users/deniz/.julia/packages/IJulia/F1GUo/src/execute_request.jl:67\n",
      "   [12] #invokelatest#1 at ./essentials.jl:790 [inlined]\n",
      "   [13] invokelatest at ./essentials.jl:789 [inlined]\n",
      "   [14] eventloop(::ZMQ.Socket) at /Users/deniz/.julia/packages/IJulia/F1GUo/src/eventloop.jl:8\n",
      "   [15] (::getfield(IJulia, Symbol(\"##15#18\")))() at ./task.jl:268\n",
      "  \n",
      "\u001b[37m\u001b[1mTest Summary:     | \u001b[22m\u001b[39m\u001b[91m\u001b[1mError  \u001b[22m\u001b[39m\u001b[36m\u001b[1mTotal\u001b[22m\u001b[39m\n",
      "Testing attention | \u001b[91m    1  \u001b[39m\u001b[36m    1\u001b[39m\n"
     ]
    },
    {
     "ename": "TestSetException",
     "evalue": "Some tests did not pass: 0 passed, 0 failed, 1 errored, 0 broken.",
     "output_type": "error",
     "traceback": [
      "Some tests did not pass: 0 passed, 0 failed, 1 errored, 0 broken.",
      "",
      "Stacktrace:",
      " [1] finish(::Test.DefaultTestSet) at /Users/sabae/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.2/Test/src/Test.jl:877",
      " [2] top-level scope at /Users/sabae/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.2/Test/src/Test.jl:1123",
      " [3] top-level scope at In[101]:2"
     ]
    }
   ],
   "source": [
    "@testset \"Testing attention\" begin\n",
    "    src1,tgt1 = first(dtrn)\n",
    "    key1,val1 = encode(pretrained, src1)\n",
    "    H,B = pretrained.encoder.hiddenSize, size(src1,1)\n",
    "    Knet.seed!(1)\n",
    "    x = KnetArray(randn(Float32,H,B,5))\n",
    "    y = pretrained.attention(x, (key1, val1))\n",
    "    @test size(y) == size(x)\n",
    "    @test norm(y) ≈ 808.381f0\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 5. Decoder\n",
    "\n",
    "`decode()` takes a model `s`, a target language minibatch `tgt`, the memory from the\n",
    "encoder `mem` and the decoder output from the previous time step `prev`. After the input\n",
    "is passed through the embedding layer, it is concatenated with `prev` (this is called\n",
    "input feeding). The resulting tensor is passed through `s.decoder`. Finally the\n",
    "`s.attention` layer takes the decoder output and the encoder memory to compute the\n",
    "\"attention vector\" which is returned by `decode()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "decode (generic function with 1 method)"
      ]
     },
     "execution_count": 102,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function decode(s::S2S, tgt, mem, prev)\n",
    "    \n",
    "    tgt_embed_tensor = dropout(s.tgtembed(tgt), s.dropout)\n",
    "    input = vcat(tgt_embed_tensor,prev)\n",
    "    y_dec = s.decoder(input)\n",
    "    s.attention(y_dec, mem)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37mTesting decoder: \u001b[39m\u001b[91m\u001b[1mTest Failed\u001b[22m\u001b[39m at \u001b[39m\u001b[1mIn[103]:9\u001b[22m\n",
      "  Expression: norm(cell) ≈ 131.21631f0\n",
      "   Evaluated: 132.8271f0 ≈ 131.21631f0\n",
      "Stacktrace:\n",
      " [1] top-level scope at \u001b[1mIn[103]:9\u001b[22m\n",
      " [2] top-level scope at \u001b[1m/Users/sabae/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.2/Test/src/Test.jl:1113\u001b[22m\n",
      " [3] top-level scope at \u001b[1mIn[103]:2\u001b[22m\n",
      "\u001b[37m\u001b[1mTest Summary:   | \u001b[22m\u001b[39m\u001b[32m\u001b[1mPass  \u001b[22m\u001b[39m\u001b[91m\u001b[1mFail  \u001b[22m\u001b[39m\u001b[36m\u001b[1mTotal\u001b[22m\u001b[39m\n",
      "Testing decoder | \u001b[32m   1  \u001b[39m\u001b[91m   1  \u001b[39m\u001b[36m    2\u001b[39m\n"
     ]
    },
    {
     "ename": "TestSetException",
     "evalue": "Some tests did not pass: 1 passed, 1 failed, 0 errored, 0 broken.",
     "output_type": "error",
     "traceback": [
      "Some tests did not pass: 1 passed, 1 failed, 0 errored, 0 broken.",
      "",
      "Stacktrace:",
      " [1] finish(::Test.DefaultTestSet) at /Users/sabae/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.2/Test/src/Test.jl:877",
      " [2] top-level scope at /Users/sabae/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.2/Test/src/Test.jl:1123",
      " [3] top-level scope at In[103]:2"
     ]
    }
   ],
   "source": [
    "@testset \"Testing decoder\" begin\n",
    "    src1,tgt1 = first(dtrn)\n",
    "    key1,val1 = encode(pretrained, src1)\n",
    "    H,B = pretrained.encoder.hiddenSize, size(src1,1)\n",
    "    Knet.seed!(1)\n",
    "    cell = randn!(similar(key1, size(key1,1), size(key1,3), 1))\n",
    "    cell = decode(pretrained, tgt1[:,1:1], (key1,val1), cell)\n",
    "    @test size(cell) == (H,B,1)\n",
    "    @test norm(cell) ≈ 131.21631f0\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 6. Loss\n",
    "\n",
    "The loss function takes source language minibatch `src`, and a target language minibatch\n",
    "`tgt` and returns `sumloss/numwords` if `average=true` or `(sumloss,numwords)` if\n",
    "`average=false` where `sumloss` is the total negative log likelihood loss and `numwords` is\n",
    "the number of words predicted (including a final eos for each sentence). The source is first\n",
    "encoded using `encode` yielding a `(keys,vals)` pair (memory). Then the decoder is called to\n",
    "predict each word of `tgt` given the previous word, `(keys,vals)` pair, and the previous\n",
    "decoder output. The previous decoder output is initialized with zeros for the first\n",
    "step. The output of the decoder at each step is passed through the projection layer giving\n",
    "word scores. Losses can be computed from word scores and masked/shifted `tgt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "function (s::S2S)(src, tgt; average=true)\n",
    "    batchsize = size(tgt,1)\n",
    "    \n",
    "    mem = encode(s, src)\n",
    "    \n",
    "    prev = zeros(Float32, size(s.projection.w, 2), batchsize, 1)\n",
    "    \n",
    "    if(gpu()>=0)\n",
    "        prev = KnetArray(prev)\n",
    "    end\n",
    "    \n",
    "    output = copy(prev)\n",
    "    \n",
    "    for i = 1:size(tgt,2)-1\n",
    "        tmp_tgt = reshape(tgt[:,i], (size(tgt[:,i], 1), 1))\n",
    "        y_dec = decode(s, tmp_tgt, mem, prev)\n",
    "        prev = y_dec\n",
    "        output = cat(output, y_dec, dims = 3)\n",
    "    end\n",
    "\n",
    "    output = output[:,:,2:end]\n",
    "    hy, b ,ty = size(output)\n",
    "    \n",
    "    output = reshape(output, (hy, b*ty))\n",
    "    \n",
    "    scores = s.projection(output)\n",
    "    y_gold = mask!(tgt[:,2:end], s.tgtvocab.eos)\n",
    "    \n",
    "    nll(scores, y_gold; average = average)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37m\u001b[1mTest Summary: | \u001b[22m\u001b[39m\u001b[32m\u001b[1mPass  \u001b[22m\u001b[39m\u001b[36m\u001b[1mTotal\u001b[22m\u001b[39m\n",
      "Testing loss  | \u001b[32m   3  \u001b[39m\u001b[36m    3\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Test.DefaultTestSet(\"Testing loss\", Any[], 3, false)"
      ]
     },
     "execution_count": 109,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@testset \"Testing loss\" begin\n",
    "    src1,tgt1 = first(dtrn)\n",
    "    @test pretrained(src1,tgt1) ≈ 1.4666592f0\n",
    "    @test pretrained(src1,tgt1,average=false)[2] == (1949.1901f0, 1329)[2]\n",
    "    @test pretrained(src1,tgt1,average=false)[1] ≈ (1949.1901f0, 1329)[1] #converted loss to similarity\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 7. Greedy translator\n",
    "\n",
    "An `S2S` object can be called with a single argument (source language minibatch `src`, with\n",
    "size `B,Tx`) to generate translations (target language minibatch with size `B,Ty`). The\n",
    "keyword argument `stopfactor` determines how much longer the output can be compared to the\n",
    "input. Similar to the loss function, the source minibatch is encoded yield a `(keys,vals)`\n",
    "pair (memory). We generate the output one time step at a time by calling the decoder with\n",
    "the last output, the memory, and the last decoder state. The last output is initialized to\n",
    "an array of `eos` tokens and the last decoder state is initialized to an array of\n",
    "zeros. After computing the scores for the next word using the projection layer, the highest\n",
    "scoring words are selected and appended to the output. The generation stops when all outputs\n",
    "in the batch have generated `eos` or when the length of the output is `stopfactor` times the\n",
    "input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {},
   "outputs": [],
   "source": [
    "function (s::S2S)(src; stopfactor = 3)\n",
    "    \n",
    "    \n",
    "    isDone = false\n",
    "    batch_size = size(src,1)\n",
    "    input = repeat([s.tgtvocab.eos], batch_size)\n",
    "    is_all_finished = zeros(batch_size)\n",
    "    translated_sentences = copy(input)\n",
    "    max_length_output = 0\n",
    "    \n",
    "    mem = encode(s, src)\n",
    "    \n",
    "    prev_decoder_output = zeros(Float32, size(s.encoder.h, 1), batch_size, 1)\n",
    "    if (gpu() >= 0)\n",
    "        prev_decoder_output = KnetArray(prev_decoder_output)\n",
    "    end\n",
    "    input = reshape(input, (length(input), 1))\n",
    "    \n",
    "    while (!isDone && max_length_output < stopfactor*size(src,2))        \n",
    "        \n",
    "        \n",
    "        y = decode(s, input, mem, prev_decoder_output)\n",
    "        prev_decoder_output = y\n",
    "        \n",
    "          \n",
    "        hy, b ,ty = size(y)\n",
    "        y = reshape(y, (hy, b*ty))\n",
    "        \n",
    "        scores = s.projection(y)\n",
    "        \n",
    "        output_words = reshape(map(x->x[1], argmax(scores, dims = 1)), batch_size)\n",
    "        translated_sentences = hcat(translated_sentences, output_words)\n",
    "       \n",
    "        max_length_output = size(translated_sentences, 2)\n",
    "        input = reshape(output_words, (length(output_words), 1))\n",
    "        \n",
    "       \n",
    "        tmp_output_words = copy(output_words)\n",
    "        tmp_output_words = tmp_output_words .== s.tgtvocab.eos\n",
    "        is_all_finished += tmp_output_words\n",
    "        if(sum(is_all_finished.==0)==0)\n",
    "            isDone = true\n",
    "        end\n",
    "    end\n",
    "    return translated_sentences[:, 2:end]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37m\u001b[1mTest Summary:      | \u001b[22m\u001b[39m\u001b[32m\u001b[1mPass  \u001b[22m\u001b[39m\u001b[36m\u001b[1mTotal\u001b[22m\u001b[39m\n",
      "Testing translator | \u001b[32m   2  \u001b[39m\u001b[36m    2\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Test.DefaultTestSet(\"Testing translator\", Any[], 2, false)"
      ]
     },
     "execution_count": 134,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@testset \"Testing translator\" begin\n",
    "    src1,tgt1 = first(dtrn)\n",
    "    tgt2 = pretrained(src1)\n",
    "    @test size(tgt2) == (64, 41)\n",
    "    @test tgt2[1:3,1:3] == [14 25 10647; 37 25 1426; 27 5 349]\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 8. Training\n",
    "\n",
    "`trainmodel` creates, trains and returns an `S2S` model. The arguments are described in\n",
    "comments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "trainmodel (generic function with 1 method)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function trainmodel(trn,                  # Training data\n",
    "                    dev,                  # Validation data, used to determine the best model\n",
    "                    tst...;               # Zero or more test datasets, their loss will be periodically reported\n",
    "                    bidirectional = true, # Whether to use a bidirectional encoder\n",
    "                    layers = 2,           # Number of layers (use `layers÷2` for a bidirectional encoder)\n",
    "                    hidden = 512,         # Size of the hidden vectors\n",
    "                    srcembed = 512,       # Size of the source language embedding vectors\n",
    "                    tgtembed = 512,       # Size of the target language embedding vectors\n",
    "                    dropout = 0.2,        # Dropout probability\n",
    "                    epochs = 0,           # Number of epochs (one of epochs or iters should be nonzero for training)\n",
    "                    iters = 0,            # Number of iterations (one of epochs or iters should be nonzero for training)\n",
    "                    bleu = false,         # Whether to calculate the BLEU score for the final model\n",
    "                    save = false,         # Whether to save the final model\n",
    "                    seconds = 60,         # Frequency of progress reporting\n",
    "                    )\n",
    "    @show bidirectional, layers, hidden, srcembed, tgtembed, dropout, epochs, iters, bleu, save; flush(stdout)\n",
    "    model = S2S(hidden, srcembed, tgtembed, trn.src.vocab, trn.tgt.vocab;\n",
    "                layers=layers, dropout=dropout, bidirectional=bidirectional)\n",
    "    \n",
    "    epochs == iters == 0 && return model\n",
    "\n",
    "    (ctrn,cdev,ctst) = collect(trn),collect(dev),collect.(tst)\n",
    "    traindata = (epochs > 0\n",
    "                 ? collect(flatten(shuffle!(ctrn) for i in 1:epochs))\n",
    "                 : shuffle!(collect(take(cycle(ctrn), iters))))\n",
    "\n",
    "    bestloss, bestmodel = loss(model, cdev), deepcopy(model)\n",
    "    progress!(adam(model, traindata), seconds=seconds) do y\n",
    "        devloss = loss(model, cdev)\n",
    "        tstloss = map(d->loss(model,d), ctst)\n",
    "        if devloss < bestloss\n",
    "            bestloss, bestmodel = devloss, deepcopy(model)\n",
    "        end\n",
    "        println(stderr)\n",
    "        (dev=devloss, tst=tstloss, mem=Float32(CuArrays.usage[]))\n",
    "    end\n",
    "    save && Knet.save(\"attn-$(Int(time_ns())).jld2\", \"model\", bestmodel)\n",
    "    bleu && Main.bleu(bestmodel,dev)\n",
    "    return bestmodel\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train a model: If your implementation is correct, the first epoch should take about 24\n",
    "minutes on a v100 and bring the loss from 9.83 to under 4.0. 10 epochs would take about 4\n",
    "hours on a v100. With other GPUs you may have to use a smaller batch size (if memory is\n",
    "lower) and longer time (if gpu speed is lower)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "S2S(Embed(P(Array{Float32,2}(512,38126))), LSTM(input=512,hidden=512,bidirectional,dropout=0.2), Memory(P(Array{Float32,2}(512,1024))), Embed(P(Array{Float32,2}(512,18857))), LSTM(input=1024,hidden=512,layers=2,dropout=0.2), Attention(1, P(Array{Float32,2}(512,1536)), P(Array{Float32,1}(1))), Linear(P(Array{Float32,2}(18857,512)), P(Array{Float32,1}(18857))), 0.2, Vocab(Dict(\"dev\" => 1277,\"komuta\" => 13566,\"ellisi\" => 25239,\"adresini\" => 22820,\"yüzeyi\" => 4051,\"paris'te\" => 9494,\"kafamdaki\" => 18790,\"yüzeyinde\" => 5042,\"geçerlidir\" => 6612,\"kökten\" => 7774…), [\"<s>\", \"<unk>\", \".\", \",\", \"bir\", \"ve\", \"bu\", \"''\", \"``\", \"için\"  …  \"seçmemiz\", \"destekleyip\", \"karşılaştırılabilir\", \"ördeğin\", \"gününüzü\", \"bağışçı\", \"istismara\", \"yaşça\", \"tedci\", \"fakültesi'nde\"], 2, 1, split), Vocab(Dict(\"middle-income\" => 13398,\"photosynthesis\" => 7689,\"polarizing\" => 17881,\"henry\" => 4248,\"abducted\" => 15691,\"rises\" => 6225,\"hampshire\" => 13888,\"whiz\" => 16835,\"cost-benefit\" => 13137,\"progression\" => 5549…), [\"<s>\", \"<unk>\", \",\", \".\", \"the\", \"and\", \"to\", \"of\", \"a\", \"that\"  …  \"archaea\", \"handshake\", \"brit\", \"wiper\", \"heroines\", \"coca\", \"exceptionally\", \"gallbladder\", \"autopsies\", \"linguistics\"], 2, 1, split))"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Uncomment the appropriate option for training:\n",
    "\n",
    "#model = pretrained  # Use reference model\n",
    "#model = Knet.load(\"attn-2888149734332.jld2\", \"model\")  # Load pretrained model\n",
    "model = trainmodel(dtrn,ddev,take(dtrn,20); epochs=10, save=true, bleu=true)  # Train model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Code to sample translations from a dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "translate_sample (generic function with 1 method)"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data1 = MTData(tr_dev, en_dev, batchsize=1) |> collect;\n",
    "function translate_sample(model, data)\n",
    "    (src,tgt) = rand(data)\n",
    "    out = model(src)\n",
    "    println(\"SRC: \", int2str(src,model.srcvocab))\n",
    "    println(\"REF: \", int2str(tgt,model.tgtvocab))\n",
    "    println(\"OUT: \", int2str(out,model.tgtvocab))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate translations for random instances from the dev set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[11, 351, 3, 5, 4069, 3, 11, 22, 195, 12, 5, 2, 3, 6, 11, 34, 40, 237, 38, 12, 678, 3, 6, 11, 34, 40, 83, 5, 275, 8, 678, 6, 11, 34, 40, 4, 1]\n",
      "[\"i\", \"remember\", \",\", \"the\", \"taliban\", \",\", \"i\", \"was\", \"still\", \"in\", \"the\", \"<unk>\", \",\", \"and\", \"i\", \"can\", \"not\", \"always\", \"be\", \"in\", \"fear\", \",\", \"and\", \"i\", \"can\", \"not\", \"get\", \"the\", \"future\", \"of\", \"fear\", \"and\", \"i\", \"can\", \"not\", \".\", \"<s>\"]\n",
      "SRC: taliban yılları boyunca hatırlıyorum bazen <unk> hayatımızdan ve hep korku içinde olmaktan ve geleceği <unk> .\n",
      "REF: during taliban years , i remember there were times i would get so frustrated by our life and always being scared and not seeing a future .\n",
      "OUT: i remember , the taliban , i was still in the <unk> , and i can not always be in fear , and i can not get the future of fear and i can not .\n"
     ]
    }
   ],
   "source": [
    "translate_sample(model, data1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Code to generate translations from user input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "translate_input (generic function with 1 method)"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function translate_input(model)\n",
    "    v = model.srcvocab\n",
    "    src = [ get(v.w2i, w, v.unk) for w in v.tokenizer(readline()) ]'\n",
    "    out = model(src)\n",
    "    println(\"SRC: \", int2str(src,model.srcvocab))\n",
    "    println(\"OUT: \", int2str(out,model.tgtvocab))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate translations for user input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# translate_input(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Competition\n",
    "\n",
    "The reference model `pretrained` has 16.2 bleu. By playing with the optimization algorithm\n",
    "and hyperparameters, using per-sentence loss, and (most importantly) splitting the Turkish\n",
    "words I was able to push the performance to 21.0 bleu. I will give extra credit to groups\n",
    "that can exceed 21.0 bleu in this dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
   ]
  }
 ],
 "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
}