Chapter15 - Intro to Federated Learning - Deep Learning on Unseen Data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "init_hidden (generic function with 2 methods)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Random\n",
    "import Base:+,-,*,println, sum, broadcasted, size, adjoint, show, dropdims, tanh, reshape\n",
    "using Base.Iterators:partition, flatten\n",
    "\n",
    "mutable struct Tensor\n",
    "    data\n",
    "    autograd\n",
    "    creators\n",
    "    creation_op\n",
    "    id\n",
    "    children\n",
    "    grad \n",
    "    index_select_indices\n",
    "    softmax_output\n",
    "    target_dist\n",
    "    parent_shape\n",
    "    \n",
    "    function Tensor(data; autograd=false, creators=nothing, creation_op = nothing, id=nothing)\n",
    "        if isnothing(id)\n",
    "            id = rand(1:100000)\n",
    "        end\n",
    "        T = new(data, autograd, creators, creation_op, id)\n",
    "        T.children = Dict()\n",
    "        T.grad = nothing\n",
    "        T.index_select_indices = nothing\n",
    "        \n",
    "        if !(isnothing(creators))\n",
    "            for c in creators\n",
    "                if haskey(c.children, T.id)\n",
    "                    c.children[T.id] += 1\n",
    "                else\n",
    "                    c.children[T.id] = 1\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "        return T\n",
    "    end\n",
    "end\n",
    "\n",
    "function all_children_grads_accounted_for(t::Tensor)\n",
    "    for (id, cnt) in t.children\n",
    "        if (cnt != 0)\n",
    "            return false\n",
    "        end\n",
    "    end\n",
    "    return true\n",
    "end\n",
    "\n",
    "function backward(t::Tensor, grad=nothing, grad_origin=nothing)\n",
    "    if t.autograd\n",
    "        if isnothing(grad)\n",
    "            grad = Tensor(ones(size(t.data)))\n",
    "        end\n",
    "    \n",
    "        if !(isnothing(grad_origin))\n",
    "            if t.children[grad_origin.id] == 0\n",
    "                return\n",
    "                throw(\"cannot backprop more than once\")\n",
    "            else\n",
    "                t.children[grad_origin.id] -= 1\n",
    "            end\n",
    "        end\n",
    "        \n",
    "        if isnothing(t.grad)\n",
    "            t.grad = grad\n",
    "        else\n",
    "            t.grad += grad\n",
    "        end\n",
    "        \n",
    "        # grads must not have grads of their own\n",
    "        @assert !grad.autograd\n",
    "        \n",
    "        # only continue backpropping if there's something to\n",
    "        # backprop into and if all gradients (from children)\n",
    "        # are accounted for override waiting for children if\n",
    "        # \"backprop\" was called on this variable directly\n",
    "        \n",
    "        if (!isnothing(t.creators) && (all_children_grads_accounted_for(t) || isnothing(grad_origin)))\n",
    "            if t.creation_op == \"add\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"sub\"\n",
    "                backward(t.creators[1], t.grad, t)\n",
    "                backward(t.creators[2], -t.grad, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mul\"\n",
    "                new_ = t.grad .* t.creators[2]\n",
    "                backward(t.creators[1], new_, t)\n",
    "                new_ = t.grad .* t.creators[1]\n",
    "                backward(t.creators[2], new_, t)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"mm\"\n",
    "                c1 = t.creators[1]\n",
    "                c2 = t.creators[2]\n",
    "                new_ =  t.grad * c2' ################\n",
    "                backward(c1, new_)\n",
    "                new_ = c1' * t.grad\n",
    "                backward(c2, new_)\n",
    "            end\n",
    "                  \n",
    "            if t.creation_op == \"transpose\"\n",
    "                backward(t.creators[1], t.grad')\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"reshape\"\n",
    "                backward(t.creators[1], reshape(t.grad, t.parent_shape))\n",
    "            end\n",
    "            \n",
    "            if occursin(\"sum\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                backward(t.creators[1], expand(t.grad, dim, size(t.creators[1].data)[dim]))\n",
    "            end\n",
    "            \n",
    "            if occursin(\"expand\", t.creation_op)\n",
    "                dim = parse(Int, split(t.creation_op, \"_\")[2])\n",
    "                ndims_cr = ndims(t.creators[1].data)\n",
    "                backward(t.creators[1], dropdims(sum(t.grad;dims=dim);dims=dim, ndims_cr=ndims_cr))\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"neg\"\n",
    "                backward(t.creators[1], -t.grad)\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"sigmoid\"\n",
    "                ones_ = Tensor(ones(size(t.grad.data)))\n",
    "                backward(t.creators[1], t.grad .* t .* (ones_ - t) )\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"tanh\"\n",
    "                ones_ = Tensor(ones(size(t.grad.data)))\n",
    "                backward(t.creators[1], t.grad .* (ones_ - (t .* t)))\n",
    "            end\n",
    "            \n",
    "            if t.creation_op == \"index_select\"\n",
    "                new_grad = zeros(size(t.creators[1]))\n",
    "                indices = t.index_select_indices.data\n",
    "                major_chunks = partition(1:size(t.grad,2),length(indices))\n",
    "                grad_chunks = [t.grad.data[:,inds][:,j]  for(i,inds) in enumerate(major_chunks) for j=1:size(inds)[1]]\n",
    "    \n",
    "                for (i,ind) in enumerate(flatten(indices))\n",
    "                    new_grad[:,ind] +=  grad_chunks[i]\n",
    "                end\n",
    "                backward(t.creators[1], Tensor(new_grad))\n",
    "            end\n",
    "            if t.creation_op == \"cross_entropy\"\n",
    "                dx = t.softmax_output .- t.target_dist\n",
    "                backward(t.creators[1], Tensor(dx))\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end\n",
    "                        \n",
    "size(a::Tensor) = size(a.data)\n",
    "size(a::Tensor, ind::Int) = size(a.data, ind)\n",
    "\n",
    "function +(a::Tensor, b::Tensor)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor(a.data + b.data; autograd=true, creators=[a,b], creation_op = \"add\")\n",
    "    end\n",
    "    return Tensor(a.data+b.data)\n",
    "end\n",
    "\n",
    "function -(a::Tensor)\n",
    "    if (a.autograd)\n",
    "        return Tensor(a.data .* -1; autograd=true, creators=[a], creation_op = \"neg\")\n",
    "    end\n",
    "    return Tensor(a.data .* -1)\n",
    "end\n",
    "\n",
    "function -(a::Tensor, b::Tensor)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor(a.data - b.data; autograd=true, creators=[a,b], creation_op = \"sub\")\n",
    "    end\n",
    "    return Tensor(a.data-b.data)\n",
    "end\n",
    "\n",
    "#element-wise multiplication\n",
    "function broadcasted(f::typeof(*), a::Tensor, b::Tensor)\n",
    "    new_data = zeros(size(a.data))\n",
    "    for i=1:length(new_data)\n",
    "        new_data[i] = f(a.data[i] ,b.data[i])\n",
    "    end\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor(new_data; autograd=true, creators=[a,b], creation_op =\"mul\")\n",
    "    end\n",
    "    return Tensor(new_data)\n",
    "end\n",
    "\n",
    "function broadcasted(f::typeof(-), a::Tensor, b::Tensor)\n",
    "    new_data = zeros(size(a.data))\n",
    "    for i=1:length(new_data)\n",
    "        new_data[i] = -(a.data[i] ,b.data[i])\n",
    "    end\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor(new_data; autograd=true, creators=[a,b], creation_op =\"sub\")\n",
    "    end\n",
    "    return Tensor(new_data)\n",
    "end\n",
    "\n",
    "function sum(a::Tensor; dims=dims)\n",
    "    new_ = dropdims(sum(a.data ;dims=dims), dims = tuple(findall(size(a) .== 1)...))\n",
    "    if (a.autograd)\n",
    "        return Tensor(new_; autograd=true, creators=[a], creation_op = \"sum_\"*string(dims))\n",
    "    end\n",
    "    return Tensor(new_)\n",
    "end\n",
    "\n",
    "function dropdims(a::Tensor;dims=dims,ndims_cr=ndims_cr)\n",
    "    if ndims(a.data) == ndims_cr\n",
    "        return a\n",
    "    end\n",
    "    if (a.autograd)\n",
    "        return Tensor(dropdims(a.data ;dims=dims); autograd=true, creators=[a], creation_op = \"dropdims\")\n",
    "    end\n",
    "    return Tensor(dropdims(a.data ;dims=dims))\n",
    "end\n",
    "\n",
    "function expand(a::Tensor, dim, copies)\n",
    "    sz = size(a)\n",
    "    rep = ntuple(d->d==dim ? copies : 1, length(sz)+1)\n",
    "    new_size = ntuple(d->d<dim ? sz[d] : d == dim ? 1 : sz[d-1], length(sz)+1)\n",
    "    new_data =  repeat(reshape(a.data, new_size), outer=rep)\n",
    "    if (a.autograd)\n",
    "        return Tensor(new_data; autograd=true, creators=[a], creation_op = \"expand_\"*string(dim))\n",
    "    end\n",
    "    return Tensor(new_data)\n",
    "end\n",
    "\n",
    "#transpose\n",
    "function adjoint(a::Tensor)\n",
    "    if (a.autograd)\n",
    "        return Tensor(a.data';autograd=true, creators=[a], creation_op = \"transpose\")\n",
    "    end\n",
    "    return Tensor(a.data')\n",
    "end\n",
    "\n",
    "#matrix multiply \n",
    "function *(a::Tensor, b::Tensor)\n",
    "    if (a.autograd && b.autograd)\n",
    "        return Tensor(a.data * b.data; autograd=true, creators=[a,b], creation_op = \"mm\")\n",
    "    end\n",
    "    return Tensor(a.data * b.data)\n",
    "end\n",
    "\n",
    "\n",
    "function index_select_helper(a::Array, indices)\n",
    "    return reduce(hcat,map(ind -> a[:,ind], indices))\n",
    "end\n",
    "\n",
    "function index_select(a::Tensor, indices::Tensor)\n",
    "    new_ = index_select_helper(a.data, indices.data)\n",
    "    if (a.autograd)\n",
    "        T = Tensor(new_, autograd=true, creators=[a], creation_op = \"index_select\")\n",
    "        T.index_select_indices = indices\n",
    "        return T\n",
    "    end\n",
    "    return Tensor(new_)\n",
    "end\n",
    "\n",
    "function reshape(a::Tensor, shape)\n",
    "    if (a.autograd)\n",
    "        T= Tensor(reshape(a.data, shape); autograd=true, creators=[a], creation_op = \"reshape\")\n",
    "        T.parent_shape = size(a.data)\n",
    "        return T\n",
    "    end\n",
    "    return Tensor(reshape(a.data, shape))\n",
    "end\n",
    "\n",
    "\n",
    "println(t::Tensor) = println(t.data)\n",
    "show(io::IO,m::MIME\"text/plain\",a::Tensor) = show(io,m,a.data)\n",
    "                        \n",
    "abstract type Layer end\n",
    "\n",
    "function get_parameters(l::Layer)\n",
    "    return l.parameters\n",
    "end\n",
    "\n",
    "mutable struct Linear <: Layer\n",
    "    W\n",
    "    b\n",
    "    use_bias\n",
    "    parameters\n",
    "                            \n",
    "    function Linear(n_inputs, n_outputs;bias=true)\n",
    "        linear = new()\n",
    "        linear.use_bias = bias\n",
    "        linear.W = Tensor(randn(n_outputs, n_inputs) .* sqrt(2.0/n_inputs), autograd=true)\n",
    "        if bias\n",
    "            linear.b = Tensor(zeros(n_outputs), autograd=true) \n",
    "            linear.parameters = [linear.W,linear.b]\n",
    "        else\n",
    "            linear.parameters = [linear.W]\n",
    "        end\n",
    "        return linear\n",
    "    end\n",
    "end\n",
    "\n",
    "function forward(l::Linear, input)\n",
    "    if l.use_bias\n",
    "        return (l.W * input)  + expand(l.b,2,size(input.data, 2))\n",
    "    end\n",
    "    return l.W * input\n",
    "end                        \n",
    "\n",
    "                        \n",
    "mutable struct Sequential <: Layer\n",
    "    layers\n",
    "    function Sequential(layers)\n",
    "        return new(layers)\n",
    "    end\n",
    "end\n",
    "\n",
    "function add(s::Sequential, layer)\n",
    "    push!(s.layers, layer)\n",
    "end\n",
    "\n",
    "function forward(s::Sequential, input)\n",
    "    for layer in s.layers\n",
    "        input = forward(layer, input)\n",
    "    end\n",
    "    return input\n",
    "end\n",
    "\n",
    "function get_parameters(s::Sequential)\n",
    "    parameters = [get_parameters(layer) for layer in s.layers]\n",
    "    return collect(Iterators.flatten(parameters))\n",
    "end\n",
    "\n",
    "mutable struct SGD\n",
    "    parameters\n",
    "    alpha\n",
    "    SGD(parameters, alpha) = new(parameters, alpha)\n",
    "end\n",
    "\n",
    "function zero!(opt::SGD)\n",
    "    for p in opt.parameters\n",
    "        p.grad.data .*= 0.0\n",
    "    end\n",
    "end\n",
    "\n",
    "function step(opt::SGD, zero=true)\n",
    "    for p in opt.parameters\n",
    "        p.data -= (p.grad.data .* opt.alpha)\n",
    "        if zero\n",
    "            p.grad.data .*= 0.0\n",
    "        end\n",
    "    end\n",
    "end\n",
    "                        \n",
    "σ(x) = 1/(1+exp(-x))                        \n",
    "\n",
    "struct Tanh <: Layer\n",
    "    Tanh() = new()\n",
    "end\n",
    "\n",
    "struct Sigmoid <: Layer\n",
    "    Sigmoid() = new()\n",
    "end\n",
    "\n",
    "function get_parameters(l::Tanh)\n",
    "    return []\n",
    "end\n",
    "\n",
    "function get_parameters(l::Sigmoid)\n",
    "    return []\n",
    "end\n",
    "\n",
    "function forward(l::Sigmoid, a::Tensor)\n",
    "    if a.autograd\n",
    "        return Tensor(σ.(a.data); autograd=true, creators=[a], creation_op = \"sigmoid\")\n",
    "    end\n",
    "    return Tensor(σ.(a.data))\n",
    "end\n",
    "        \n",
    "function forward(l::Tanh, a::Tensor)\n",
    "    if a.autograd\n",
    "        return Tensor(tanh.(a.data); autograd=true, creators=[a], creation_op = \"tanh\")\n",
    "    end\n",
    "    return Tensor(tanh.(a.data))\n",
    "end    \n",
    "                        \n",
    "                        \n",
    "mutable struct Embedding <: Layer\n",
    "    vocab_size\n",
    "    dim\n",
    "    weight\n",
    "    parameters\n",
    "    # this random initialiation style is just a convention from word2vec\n",
    "    function Embedding(dim, vocab_size) \n",
    "        E = new(vocab_size, dim, Tensor((randn(dim, vocab_size) .- 0.5) ./ dim; autograd=true))\n",
    "        E.parameters = [E.weight]\n",
    "        return E\n",
    "    end\n",
    "end\n",
    "\n",
    "function forward(E::Embedding, indices)\n",
    "    return index_select(E.weight, indices)\n",
    "end\n",
    "                        \n",
    "using Statistics: mean\n",
    "using LinearAlgebra: I\n",
    "function softmax(x)\n",
    "    temp = exp.(x)\n",
    "    return temp ./ sum(temp;dims=1)\n",
    "end\n",
    "\n",
    "struct MSELoss <: Layer\n",
    "    MSELoss() = new()\n",
    "end\n",
    "\n",
    "function forward(l::MSELoss, pred, target)\n",
    "    return sum((pred - target) .* (pred - target);dims=2)\n",
    "end\n",
    "\n",
    "struct CrossEntropyLoss \n",
    "    CrossEntropyLoss() = new()\n",
    "end\n",
    "\n",
    "function forward(l::CrossEntropyLoss, a::Tensor, target::Tensor)\n",
    "    softmax_output = softmax(a.data)\n",
    "    log_out = log.(softmax_output)\n",
    "    sz = size(a.data, 1)\n",
    "    identity = 1.0 .* Matrix(I, (sz, sz))\n",
    "    target_dist = reshape(identity[:,target.data],(size(a.data)))\n",
    "    loss = -mean(sum(log_out .* target_dist;dims=1))\n",
    "    if a.autograd\n",
    "        loss = Tensor(loss; autograd=true, creators=[a], creation_op = \"cross_entropy\")\n",
    "        loss.softmax_output = softmax_output\n",
    "        loss.target_dist = target_dist\n",
    "        return loss\n",
    "    end\n",
    "    return Tensor(loss)\n",
    "end\n",
    "\n",
    "\n",
    "mutable struct RNNCell <: Layer\n",
    "    n_hidden\n",
    "    \n",
    "    activation\n",
    "    \n",
    "    w_ih\n",
    "    w_hh\n",
    "    w_ho\n",
    "    \n",
    "    parameters\n",
    "    \n",
    "    function RNNCell(n_inputs, n_hidden, n_output, activation=\"sigmoid\")\n",
    "        if activation == \"sigmoid\"\n",
    "            act = Sigmoid()\n",
    "        elseif activation == \"tanh\"\n",
    "            act = Tanh()\n",
    "        else\n",
    "            throw(\"Non-linearity not found\")\n",
    "        end\n",
    "        \n",
    "        parameters = []\n",
    "\n",
    "        w_ih = Linear(n_inputs, n_hidden)\n",
    "        w_hh = Linear(n_hidden, n_hidden)\n",
    "        w_ho = Linear(n_hidden, n_output)\n",
    "        \n",
    "        push!(parameters, get_parameters(w_ih))\n",
    "        push!(parameters, get_parameters(w_hh))\n",
    "        push!(parameters, get_parameters(w_ho))\n",
    "        parameters = collect(Iterators.flatten(parameters))\n",
    "        return new(n_hidden, act, w_ih, w_hh, w_ho, parameters)\n",
    "    end\n",
    "end\n",
    "\n",
    "function forward(rnn::RNNCell, input::Tensor, hidden::Tensor)\n",
    "    from_prev_hidden = forward(rnn.w_hh, hidden)\n",
    "    combined = forward(rnn.w_ih, input) + from_prev_hidden\n",
    "    new_hidden = forward(rnn.activation, combined)\n",
    "    output = forward(rnn.w_ho, new_hidden)\n",
    "    return output, new_hidden\n",
    "end\n",
    "\n",
    "function init_hidden(rnn::RNNCell; batch_size=1)\n",
    "    return Tensor(zeros(rnn.n_hidden, batch_size), autograd=true)\n",
    "end\n",
    "\n",
    "mutable struct LSTMCell <: Layer\n",
    "    \n",
    "    n_hidden\n",
    "    \n",
    "    xf\n",
    "    xi\n",
    "    xo\n",
    "    xc\n",
    "    \n",
    "    hf\n",
    "    hi\n",
    "    ho\n",
    "    hc\n",
    "    \n",
    "    w_ho\n",
    "    parameters\n",
    "    sigmoid\n",
    "    tanh\n",
    "    \n",
    "    function LSTMCell(n_inputs, n_hidden, n_output)\n",
    "\n",
    "        xf = Linear(n_inputs, n_hidden)\n",
    "        xi = Linear(n_inputs, n_hidden)\n",
    "        xo = Linear(n_inputs, n_hidden)        \n",
    "        xc = Linear(n_inputs, n_hidden) \n",
    "        \n",
    "        hf = Linear(n_hidden, n_hidden; bias=false)\n",
    "        hi = Linear(n_hidden, n_hidden; bias=false)\n",
    "        ho = Linear(n_hidden, n_hidden; bias=false)\n",
    "        hc = Linear(n_hidden, n_hidden; bias=false) \n",
    "        \n",
    "        w_ho = Linear(n_hidden, n_output; bias=false)\n",
    "        \n",
    "        parameters = [get_parameters(i) for i in [xf, xi, xo, xc, hf, hi, hc, w_ho]]\n",
    "        parameters = collect(Iterators.flatten(parameters))\n",
    "        \n",
    "        return new(n_hidden, xf, xi, xo, xc, hf, hi, ho, hc, w_ho, parameters, Sigmoid(), Tanh())\n",
    "    end\n",
    "end\n",
    "\n",
    "function forward(lstm::LSTMCell, input::Tensor, hidden)\n",
    "    \n",
    "    prev_hidden = hidden[1]        \n",
    "    prev_cell = hidden[2]\n",
    "    \n",
    "    f = forward(lstm.xf, input) + forward(lstm.sigmoid, forward(lstm.hf, prev_hidden))\n",
    "    i = forward(lstm.xi, input) + forward(lstm.sigmoid, forward(lstm.hi, prev_hidden))\n",
    "    o = forward(lstm.xo, input) + forward(lstm.sigmoid, forward(lstm.ho, prev_hidden))\n",
    "    g = forward(lstm.xc, input) + forward(lstm.tanh, forward(lstm.hc, prev_hidden))\n",
    "    \n",
    "    c = (f .* prev_cell) + (i .* g)\n",
    "\n",
    "    h = o .* forward(lstm.tanh, c)\n",
    "    \n",
    "    output = forward(lstm.w_ho, h)\n",
    "    \n",
    "    return output, (h, c)\n",
    "end\n",
    "\n",
    "function init_hidden(lstm; batch_size=1)\n",
    "    init_hidden = Tensor(zeros(lstm.n_hidden, batch_size), autograd=true)\n",
    "    init_cell   = Tensor(zeros(lstm.n_hidden, batch_size), autograd=true)\n",
    "    \n",
    "    init_hidden.data[1,:] .+= 1\n",
    "    init_cell.data[1,:] .+= 1\n",
    "    return (init_hidden, init_cell)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Plain Ole Fashioned Deep Learning (Email Spam Detection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw = readlines(\"spam.txt\")\n",
    "raw = filter.(isascii,raw)\n",
    "vocab = Set()\n",
    "\n",
    "spam = []\n",
    "\n",
    "for row in raw\n",
    "    push!(spam, Set(split(row[1:end-1])))\n",
    "    for word in spam[end]\n",
    "        push!(vocab, word)\n",
    "    end\n",
    "end\n",
    "\n",
    "raw = readlines(\"ham.txt\")\n",
    "raw = filter.(isascii,raw)\n",
    "\n",
    "ham = []\n",
    "for row in raw\n",
    "    push!(ham, Set(split(row[1:end-1])))\n",
    "    for word in ham[end]\n",
    "        push!(vocab, word)\n",
    "    end\n",
    "end\n",
    "\n",
    "push!(vocab, \"<unk>\")\n",
    "\n",
    "vocab = collect(vocab)\n",
    "\n",
    "w2i = Dict()\n",
    "for (i,w) in enumerate(vocab)\n",
    "    w2i[w] = i\n",
    "end\n",
    "\n",
    "function to_indices(input; l=500)\n",
    "    indices = []\n",
    "    for line in input\n",
    "        if length(line) < l\n",
    "            line = cat(collect(line), repeat([\"<unk>\"], l-length(line));dims=1)\n",
    "            idxs = []\n",
    "            \n",
    "            for word in line\n",
    "                \n",
    "                push!(idxs, w2i[word])\n",
    "            end\n",
    "            push!(indices, idxs)\n",
    "        end\n",
    "    end\n",
    "    return indices\n",
    "end\n",
    "\n",
    "spam_idx = to_indices(spam)\n",
    "ham_idx = to_indices(ham)\n",
    "\n",
    "train_spam_idx = spam_idx[1:end-1001]\n",
    "train_ham_idx = ham_idx[1:end-1001]\n",
    "\n",
    "test_spam_idx = spam_idx[end-1000:end]\n",
    "test_ham_idx = ham_idx[end-1000:end]\n",
    "\n",
    "train_data = []\n",
    "train_target = []\n",
    "\n",
    "test_data = []\n",
    "test_target = []\n",
    "\n",
    "for i=1:maximum([length(train_spam_idx), length(train_ham_idx)])\n",
    "    push!(train_data, train_spam_idx[((i-1)%length(train_spam_idx))+1])\n",
    "    push!(train_target,[1])\n",
    "    \n",
    "    push!(train_data, train_ham_idx[((i-1)%length(train_ham_idx))+1])\n",
    "    push!(train_target,[0])\n",
    "end\n",
    "\n",
    "for i=1:maximum([length(test_spam_idx), length(test_ham_idx)])\n",
    "    push!(test_data, test_spam_idx[((i-1)%length(test_spam_idx))+1])\n",
    "    push!(test_target,[1])\n",
    "    \n",
    "    push!(test_data, test_ham_idx[((i-1)%length(test_ham_idx))+1])\n",
    "    push!(test_target,[0])\n",
    "end\n",
    "train_target = reduce(hcat,train_target)\n",
    "test_target = reduce(hcat, test_target);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Embedding(1,length(vocab))\n",
    "sigmoid = Sigmoid()\n",
    "model.weight.data .*= 0\n",
    "criterion = MSELoss()\n",
    "optim = SGD(get_parameters(model), 0.01);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "train (generic function with 1 method)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Base.Iterators: partition\n",
    "function train(model,input_data, target_data; batch_size=500, iterations=5)\n",
    "    criterion = MSELoss()\n",
    "    optim = SGD(get_parameters(model), 0.01)\n",
    "    \n",
    "    n_batches = trunc(Int,length(input_data) / batch_size)\n",
    "    for iter=1:iterations\n",
    "        iter_loss = 0\n",
    "        for (j,i) in enumerate(partition(1:length(input_data), batch_size))\n",
    "            # padding token should stay at 0\n",
    "            model.weight.data[:,w2i[\"<unk>\"]] *= 0 \n",
    "            \n",
    "            input = Tensor(input_data[i], autograd=true)\n",
    "            target = Tensor(target_data[:,i], autograd=true)\n",
    "            embed_out = reshape(forward(model, input), (1,length(input_data[1]),length(i)))\n",
    "            pred = forward(sigmoid, sum(embed_out;dims=2))\n",
    "            \n",
    "            loss = forward(criterion, pred,target)\n",
    "            backward(loss)\n",
    "            step(optim)\n",
    "\n",
    "            iter_loss += loss.data[1] / batch_size\n",
    "            print(\"\\tLoss: $(iter_loss/(j+1)) \\r\")\n",
    "        end\n",
    "        println()\n",
    "    end\n",
    "   return model \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "test (generic function with 1 method)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function test(model, test_data, test_target)\n",
    "    model.weight.data[:,w2i[\"<unk>\"]] *= 0 \n",
    "    input = Tensor(test_data, autograd=true)\n",
    "    # target = Tensor(test_target, autograd=true)\n",
    "    embed_out = reshape(forward(model, input), (1,500,:))\n",
    "    pred = forward(sigmoid, sum(embed_out;dims=2))\n",
    "    pred_labels = Int.(pred.data .>= 0.5) .== test_target\n",
    "    return round(100*sum(pred_labels)/length(test_target);digits=2)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss: 0.036542026216407765 \n",
      "Correct on Test Set: 98.5\n",
      "Loss: 0.011293550777292613 \n",
      "Correct on Test Set: 99.15\n",
      "Loss: 0.008138803819084366 \n",
      "Correct on Test Set: 99.4\n"
     ]
    }
   ],
   "source": [
    "for i=1:3\n",
    "    model = train(model, train_data, train_target;batch_size=500,iterations=1)\n",
    "    println(\"Correct on Test Set: $(test(model, test_data, test_target))\")\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic Federated Learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "bob = (train_data[1:10000], train_target[:,1:10000])\n",
    "alice = (train_data[10001:20000], train_target[:,10001:20000])\n",
    "sue = (train_data[20001:end], train_target[:,20001:end]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Embedding(1,length(vocab))\n",
    "model.weight.data .*= 0;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting Training Round...\n",
      "\tStep 1: send the model to Bob\n",
      "\tLoss: 0.08677672659327845 \n",
      "\n",
      "\tStep 2: send the model to Alice\n",
      "\tLoss: 0.08217674407715851 \n",
      "\n",
      "\tStep 3: Send the model to Sue\n",
      "\tLoss: 0.05120595136342707  \n",
      "\n",
      "\tAverage Everyone's New Models\n",
      "Correct on Test Set: 97.65\n",
      "\n",
      "Repeat!!\n",
      "\n",
      "Starting Training Round...\n",
      "\tStep 1: send the model to Bob\n",
      "\tLoss: 0.02125951596629757  \n",
      "\n",
      "\tStep 2: send the model to Alice\n",
      "\tLoss: 0.02054173204982453  \n",
      "\n",
      "\tStep 3: Send the model to Sue\n",
      "\tLoss: 0.01952674364865583  \n",
      "\n",
      "\tAverage Everyone's New Models\n",
      "Correct on Test Set: 98.3\n",
      "\n",
      "Repeat!!\n",
      "\n",
      "Starting Training Round...\n",
      "\tStep 1: send the model to Bob\n",
      "\tLoss: 0.015777530894521444 \n",
      "\n",
      "\tStep 2: send the model to Alice\n",
      "\tLoss: 0.015155423992275244 \n",
      "\n",
      "\tStep 3: Send the model to Sue\n",
      "\tLoss: 0.014909027517198395 \n",
      "\n",
      "\tAverage Everyone's New Models\n",
      "Correct on Test Set: 98.65\n",
      "\n",
      "Repeat!!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i=1:3\n",
    "    println(\"Starting Training Round...\")\n",
    "    \n",
    "    println(\"\\tStep 1: send the model to Bob\")\n",
    "    bob_model = train(deepcopy(model), bob[1], bob[2]; iterations=1)\n",
    "    \n",
    "    println(\"\\n\\tStep 2: send the model to Alice\")\n",
    "    alice_model = train(deepcopy(model), alice[1], alice[2]; iterations=1)\n",
    "    \n",
    "    println(\"\\n\\tStep 3: Send the model to Sue\")\n",
    "    sue_model = train(deepcopy(model), sue[1], sue[2]; iterations=1)\n",
    "    \n",
    "    println(\"\\n\\tAverage Everyone's New Models\")\n",
    "    model.weight.data = (bob_model.weight.data + \n",
    "                         alice_model.weight.data + \n",
    "                         sue_model.weight.data)./3\n",
    "    \n",
    "    println(\"Correct on Test Set: $(test(model, test_data, test_target))\")\n",
    "    println(\"\\nRepeat!!\\n\")\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hacking Federated Learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tLoss: 0.125 \r\n"
     ]
    }
   ],
   "source": [
    "bobs_email = [\"my\", \"computer\", \"password\", \"is\", \"pizza\"]\n",
    "\n",
    "bob_input = [[w2i[x] for x in bobs_email], [w2i[x] for x in bobs_email]]\n",
    "bob_target = reshape([0,0],(1,2))\n",
    "\n",
    "model = Embedding(1,length(vocab))\n",
    "model.weight.data .*= 0;\n",
    "\n",
    "bobs_model = train(deepcopy(model), bob_input, bob_target; iterations=1, batch_size=2);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "password\n",
      "pizza\n",
      "is\n",
      "my\n",
      "computer\n"
     ]
    }
   ],
   "source": [
    "for (i,v) in enumerate(bobs_model.weight.data - model.weight.data)\n",
    "    if (v != 0)\n",
    "        println(vocab[i])\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Homomorphic Encryption"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Pkg\n",
    "Pkg.add(\"Paillier\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Paillier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The Answer: 8\n"
     ]
    }
   ],
   "source": [
    "public_key, private_key = Paillier.generate_paillier_keypair(1024)\n",
    "\n",
    "# encrypt the number \"5\"\n",
    "x = Paillier.encrypt(public_key, 5)\n",
    "\n",
    "# encrypt the number \"3\"\n",
    "y = Paillier.encrypt(public_key, 3)\n",
    "\n",
    "# add the two encrypted values\n",
    "z = x+y\n",
    "\n",
    "# decrypt the result\n",
    "z_ = Paillier.decrypt(private_key, z)\n",
    "println(\"The Answer: \", z_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Secure Aggregation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "train_and_encrypt (generic function with 1 method)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Embedding(1,length(vocab))\n",
    "model.weight.data .*= 0\n",
    "\n",
    "# note that in production the n_length should be at least 1024\n",
    "publickey, privatekey = Paillier.generate_paillier_keypair(128)\n",
    "encoding = Paillier.Encoding(Float64, publickey)\n",
    "\n",
    "\n",
    "function train_and_encrypt(model, input, target, encoding; iterations=1)\n",
    "    new_model = train(model, input, target; iterations=iterations)\n",
    "    encrypted_weights = Paillier.encode_and_encrypt(new_model.weight.data, encoding)\n",
    "    return encrypted_weights\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting Training Round...\n",
      "\tStep 1: send the model to Bob\n",
      "\tLoss: 0.08677672659327845 \n",
      "\n",
      "\tStep 2: send the model to Alice\n",
      "\tLoss: 0.08217674407715851 \n",
      "\n",
      "\tStep 3: Send the model to Sue\n",
      "\tLoss: 0.05120595136342707  \n",
      "\n",
      "\tStep 4: Bob, Alice, and Sue send their\tencrypted models to each other.\n",
      "Correct on Test Set: 97.65\n",
      "\n",
      "Repeat!!\n",
      "\n",
      "Starting Training Round...\n",
      "\tStep 1: send the model to Bob\n",
      "\tLoss: 0.02125951596629757  \n",
      "\n",
      "\tStep 2: send the model to Alice\n",
      "\tLoss: 0.02054173204982453  \n",
      "\n",
      "\tStep 3: Send the model to Sue\n",
      "\tLoss: 0.01952674364865583  \n",
      "\n",
      "\tStep 4: Bob, Alice, and Sue send their\tencrypted models to each other.\n",
      "Correct on Test Set: 98.3\n",
      "\n",
      "Repeat!!\n",
      "\n",
      "Starting Training Round...\n",
      "\tStep 1: send the model to Bob\n",
      "\tLoss: 0.015777530894521448 \n",
      "\n",
      "\tStep 2: send the model to Alice\n",
      "\tLoss: 0.015155423992275244 \n",
      "\n",
      "\tStep 3: Send the model to Sue\n",
      "\tLoss: 0.014909027517198394 \n",
      "\n",
      "\tStep 4: Bob, Alice, and Sue send their\tencrypted models to each other.\n",
      "Correct on Test Set: 98.65\n",
      "\n",
      "Repeat!!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i=1:3\n",
    "    println(\"Starting Training Round...\")\n",
    "    \n",
    "    println(\"\\tStep 1: send the model to Bob\")\n",
    "    bob_model = train_and_encrypt(deepcopy(model), bob[1], bob[2], encoding; iterations=1)\n",
    "    \n",
    "    \n",
    "    println(\"\\n\\tStep 2: send the model to Alice\")\n",
    "    alice_model = train_and_encrypt(deepcopy(model), alice[1], alice[2], encoding; iterations=1)\n",
    "    \n",
    "    println(\"\\n\\tStep 3: Send the model to Sue\")\n",
    "    sue_model = train_and_encrypt(deepcopy(model), sue[1], sue[2], encoding; iterations=1)\n",
    "    \n",
    "    \n",
    "    print(\"\\n\\tStep 4: Bob, Alice, and Sue send their\")\n",
    "    println(\"\\tencrypted models to each other.\")\n",
    "    \n",
    "    aggregated_model = zeros(size(bob_model))\n",
    "    for i=1:length(bob_model)\n",
    "        aggregated_model[i] = Paillier.decrypt_and_decode(privatekey, bob_model[i] + alice_model[i] + sue_model[i])/3\n",
    "    end\n",
    "    model.weight.data = aggregated_model\n",
    "    \n",
    "    println(\"Correct on Test Set: $(test(model, test_data, test_target))\")\n",
    "    println(\"\\nRepeat!!\\n\")\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.4.1",
   "language": "julia",
   "name": "julia-1.4"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.4.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}