bidaf_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "    \"\"\"\n",
    "    TODO: ADD DROPOUT\n",
    "    \"\"\"\n",
    "    def __init__(self, vocab_size, embedding_dim):\n",
    "        super(CNN, self).__init__()\n",
    "        \n",
    "        self.conv = nn.Conv1d(in_channels=embedding_dim, out_channels=100, kernel_size=5)\n",
    "        self.fc = nn.Linear(100, 100)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        #print(x.shape)\n",
    "        #conv1d wants (N, C, L)\n",
    "        x = x.permute(0, 2, 1)\n",
    "        #print(x.shape)\n",
    "        x = self.conv(x)\n",
    "        #print(x.shape)\n",
    "        x = F.max_pool1d(x, x.size()[2])\n",
    "        #print(x.shape)\n",
    "        x = x.squeeze(2)\n",
    "        x = self.fc(x)\n",
    "        #print(x.shape)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'x = Variable(torch.LongTensor(np.random.randint(0,50,size=(32,10))))\\n#50 is vocab size, 64 is embedding size\\ncnn = CNN(50, 64)\\ncnn(x)'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"x = Variable(torch.LongTensor(np.random.randint(0,50,size=(32,10))))\n",
    "#50 is vocab size, 64 is embedding size\n",
    "cnn = CNN(50, 64)\n",
    "cnn(x)\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTM(nn.Module):\n",
    "    def __init__(self, embedding_dim, hidden_dim, n_layers=1, bidirectional=True, dropout=0.2):\n",
    "        super(LSTM, self).__init__()\n",
    "        \n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers=n_layers, bidirectional=bidirectional, dropout=dropout)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        #print(x.shape) #[b, s, e]\n",
    "        x = x.permute(1, 0, 2)\n",
    "        #print(x.shape) #[s, b, e]\n",
    "        x, (h, c) = self.lstm(x)\n",
    "        #print(x.shape) #[s, b, h*2]\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Highway(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    COPIED FROM ALLENNLP GITHUB\n",
    "    \n",
    "    A `Highway layer <https://arxiv.org/abs/1505.00387>`_ does a gated combination of a linear\n",
    "    transformation and a non-linear transformation of its input.  :math:`y = g * x + (1 - g) *\n",
    "    f(A(x))`, where :math:`A` is a linear transformation, :math:`f` is an element-wise\n",
    "    non-linearity, and :math:`g` is an element-wise gate, computed as :math:`sigmoid(B(x))`.\n",
    "    This module will apply a fixed number of highway layers to its input, returning the final\n",
    "    result.\n",
    "    Parameters\n",
    "    ----------\n",
    "    input_dim : ``int``\n",
    "        The dimensionality of :math:`x`.  We assume the input has shape ``(batch_size,\n",
    "        input_dim)``.\n",
    "    num_layers : ``int``, optional (default=``1``)\n",
    "        The number of highway layers to apply to the input.\n",
    "    activation : ``Callable[[torch.Tensor], torch.Tensor]``, optional (default=``torch.nn.functional.relu``)\n",
    "        The non-linearity to use in the highway layers.\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 embedding_dim,\n",
    "                 num_layers,\n",
    "                 activation = torch.nn.functional.relu):\n",
    "        super(Highway, self).__init__()\n",
    "        self._embedding_dim = embedding_dim\n",
    "        self._layers = torch.nn.ModuleList([torch.nn.Linear(embedding_dim, embedding_dim * 2)\n",
    "                                            for _ in range(num_layers)])\n",
    "        self._activation = activation\n",
    "                \n",
    "        for layer in self._layers:\n",
    "            # We should bias the highway layer to just carry its input forward.  We do that by\n",
    "            # setting the bias on `B(x)` to be positive, because that means `g` will be biased to\n",
    "            # be high, to we will carry the input forward.  The bias on `B(x)` is the second half\n",
    "            # of the bias vector in each Linear layer.\n",
    "            layer.bias[embedding_dim:].data.fill_(1)\n",
    "\n",
    "    def forward(self, inputs: torch.Tensor) -> torch.Tensor:  # pylint: disable=arguments-differ\n",
    "        current_input = inputs\n",
    "        for layer in self._layers:\n",
    "            projected_input = layer(current_input)\n",
    "            linear_part = current_input\n",
    "            # NOTE: if you modify this, think about whether you should modify the initialization\n",
    "            # above, too.\n",
    "            nonlinear_part = projected_input[:, (0 * self._embedding_dim):(1 * self._embedding_dim)]\n",
    "            gate = projected_input[:, (1 * self._embedding_dim):(2 * self._embedding_dim)]\n",
    "            nonlinear_part = self._activation(nonlinear_part)\n",
    "            gate = torch.nn.functional.sigmoid(gate)\n",
    "            current_input = gate * linear_part + (1 - gate) * nonlinear_part\n",
    "        return current_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Embedding(nn.Module):\n",
    "    \"\"\"\n",
    "    NEED ONE FOR CHARS AND ONE FOR WORDS\"\"\"\n",
    "    def __init__(self, vocab_size, embedding_dim):\n",
    "        super(Embedding, self).__init__()\n",
    "        \n",
    "        self.embedding = nn.Embedding(vocab_size, embedding_dim)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.embedding(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MatrixAttention(nn.Module):\n",
    "    '''\n",
    "    This ``Module`` takes two matrices as input and returns a matrix of attentions.\n",
    "    We compute the similarity between each row in each matrix and return unnormalized similarity\n",
    "    scores.  Because these scores are unnormalized, we don't take a mask as input; it's up to the\n",
    "    caller to deal with masking properly when this output is used.\n",
    "    By default similarity is computed with a dot product, but you can alternatively use a\n",
    "    parameterized similarity function if you wish.\n",
    "    This is largely similar to using ``TimeDistributed(Attention)``, except the result is\n",
    "    unnormalized.  You should use this instead of ``TimeDistributed(Attention)`` if you want to\n",
    "    compute multiple normalizations of the attention matrix.\n",
    "    Input:\n",
    "        - matrix_1: ``(batch_size, num_rows_1, embedding_dim)``\n",
    "        - matrix_2: ``(batch_size, num_rows_2, embedding_dim)``\n",
    "    Output:\n",
    "        - ``(batch_size, num_rows_1, num_rows_2)``\n",
    "    Parameters\n",
    "    ----------\n",
    "    similarity_function: ``SimilarityFunction``, optional (default=``DotProductSimilarity``)\n",
    "        The similarity function to use when computing the attention.\n",
    "    '''\n",
    "    def __init__(self, similarity_function) -> None:\n",
    "        super(MatrixAttention, self).__init__()\n",
    "\n",
    "        self._similarity_function = similarity_function\n",
    "\n",
    "    def forward(self, matrix_1: torch.Tensor, matrix_2: torch.Tensor) -> torch.Tensor:\n",
    "        # pylint: disable=arguments-differ\n",
    "        tiled_matrix_1 = matrix_1.unsqueeze(2).expand(matrix_1.size()[0],\n",
    "                                                      matrix_1.size()[1],\n",
    "                                                      matrix_2.size()[1],\n",
    "                                                      matrix_1.size()[2])\n",
    "        tiled_matrix_2 = matrix_2.unsqueeze(1).expand(matrix_2.size()[0],\n",
    "                                                      matrix_1.size()[1],\n",
    "                                                      matrix_2.size()[1],\n",
    "                                                      matrix_2.size()[2])\n",
    "        return self._similarity_function(tiled_matrix_1, tiled_matrix_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearSimilarity(nn.Module):\n",
    "    \"\"\"\n",
    "    This similarity function performs a dot product between a vector of weights and some\n",
    "    combination of the two input vectors, followed by an (optional) activation function.  The\n",
    "    combination used is configurable.\n",
    "    If the two vectors are ``x`` and ``y``, we allow the following kinds of combinations: ``x``,\n",
    "    ``y``, ``x*y``, ``x+y``, ``x-y``, ``x/y``, where each of those binary operations is performed\n",
    "    elementwise.  You can list as many combinations as you want, comma separated.  For example, you\n",
    "    might give ``x,y,x*y`` as the ``combination`` parameter to this class.  The computed similarity\n",
    "    function would then be ``w^T [x; y; x*y] + b``, where ``w`` is a vector of weights, ``b`` is a\n",
    "    bias parameter, and ``[;]`` is vector concatenation.\n",
    "    Note that if you want a bilinear similarity function with a diagonal weight matrix W, where the\n",
    "    similarity function is computed as `x * w * y + b` (with `w` the diagonal of `W`), you can\n",
    "    accomplish that with this class by using \"x*y\" for `combination`.\n",
    "    Parameters\n",
    "    ----------\n",
    "    tensor_1_dim : ``int``\n",
    "        The dimension of the first tensor, ``x``, described above.  This is ``x.size()[-1]`` - the\n",
    "        length of the vector that will go into the similarity computation.  We need this so we can\n",
    "        build weight vectors correctly.\n",
    "    tensor_2_dim : ``int``\n",
    "        The dimension of the second tensor, ``y``, described above.  This is ``y.size()[-1]`` - the\n",
    "        length of the vector that will go into the similarity computation.  We need this so we can\n",
    "        build weight vectors correctly.\n",
    "    combination : ``str``, optional (default=\"x,y\")\n",
    "        Described above.\n",
    "    activation : ``Activation``, optional (default=linear (i.e. no activation))\n",
    "        An activation function applied after the ``w^T * [x;y] + b`` calculation.  Default is no\n",
    "        activation.\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 tensor_1_dim: int,\n",
    "                 tensor_2_dim: int,\n",
    "                 combination: str = 'x,y') -> None:\n",
    "        super(LinearSimilarity, self).__init__()\n",
    "        self._combination = combination\n",
    "        combined_dim = get_combined_dim(combination, [tensor_1_dim, tensor_2_dim])\n",
    "        self._weight_vector = nn.Parameter(torch.Tensor(combined_dim))\n",
    "        self._bias = nn.Parameter(torch.Tensor(1))\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        std = math.sqrt(6 / (self._weight_vector.size(0) + 1))\n",
    "        self._weight_vector.data.uniform_(-std, std)\n",
    "        self._bias.data.fill_(0)\n",
    "\n",
    "    def forward(self, tensor_1: torch.Tensor, tensor_2: torch.Tensor) -> torch.Tensor:\n",
    "        combined_tensors = combine_tensors(self._combination, [tensor_1, tensor_2])\n",
    "        dot_product = torch.matmul(combined_tensors, self._weight_vector)\n",
    "        return dot_product + self._bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def combine_tensors(combination: str, tensors) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Combines a list of tensors using element-wise operations and concatenation, specified by a\n",
    "    ``combination`` string.  The string refers to (1-indexed) positions in the input tensor list,\n",
    "    and looks like ``\"1,2,1+2,3-1\"``.\n",
    "    We allow the following kinds of combinations: ``x``, ``x*y``, ``x+y``, ``x-y``, and ``x/y``,\n",
    "    where ``x`` and ``y`` are positive integers less than or equal to ``len(tensors)``.  Each of\n",
    "    the binary operations is performed elementwise.  You can give as many combinations as you want\n",
    "    in the ``combination`` string.  For example, for the input string ``\"1,2,1*2\"``, the result\n",
    "    would be ``[1;2;1*2]``, as you would expect, where ``[;]`` is concatenation along the last\n",
    "    dimension.\n",
    "    If you have a fixed, known way to combine tensors that you use in a model, you should probably\n",
    "    just use something like ``torch.cat([x_tensor, y_tensor, x_tensor * y_tensor])``.  This\n",
    "    function adds some complexity that is only necessary if you want the specific combination used\n",
    "    to be `configurable`.\n",
    "    If you want to do any element-wise operations, the tensors involved in each element-wise\n",
    "    operation must have the same shape.\n",
    "    This function also accepts ``x`` and ``y`` in place of ``1`` and ``2`` in the combination\n",
    "    string.\n",
    "    \"\"\"\n",
    "    if len(tensors) > 9:\n",
    "        raise ConfigurationError(\"Double-digit tensor lists not currently supported\")\n",
    "    combination = combination.replace('x', '1').replace('y', '2')\n",
    "    to_concatenate = [_get_combination(piece, tensors) for piece in combination.split(',')]\n",
    "    return torch.cat(to_concatenate, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_combined_dim(combination: str, tensor_dims) -> int:\n",
    "    \"\"\"\n",
    "    For use with :func:`combine_tensors`.  This function computes the resultant dimension when\n",
    "    calling ``combine_tensors(combination, tensors)``, when the tensor dimension is known.  This is\n",
    "    necessary for knowing the sizes of weight matrices when building models that use\n",
    "    ``combine_tensors``.\n",
    "    Parameters\n",
    "    ----------\n",
    "    combination : ``str``\n",
    "        A comma-separated list of combination pieces, like ``\"1,2,1*2\"``, specified identically to\n",
    "        ``combination`` in :func:`combine_tensors`.\n",
    "    tensor_dims : ``List[int]``\n",
    "        A list of tensor dimensions, where each dimension is from the `last axis` of the tensors\n",
    "        that will be input to :func:`combine_tensors`.\n",
    "    \"\"\"\n",
    "    if len(tensor_dims) > 9:\n",
    "        raise ConfigurationError(\"Double-digit tensor lists not currently supported\")\n",
    "    combination = combination.replace('x', '1').replace('y', '2')\n",
    "    return sum([_get_combination_dim(piece, tensor_dims) for piece in combination.split(',')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_combination(combination: str, tensors) -> torch.Tensor:\n",
    "    if combination.isdigit():\n",
    "        index = int(combination) - 1\n",
    "        return tensors[index]\n",
    "    else:\n",
    "        if len(combination) != 3:\n",
    "            raise ConfigurationError(\"Invalid combination: \" + combination)\n",
    "        first_tensor = _get_combination(combination[0], tensors)\n",
    "        second_tensor = _get_combination(combination[2], tensors)\n",
    "        operation = combination[1]\n",
    "        if operation == '*':\n",
    "            return first_tensor * second_tensor\n",
    "        elif operation == '/':\n",
    "            return first_tensor / second_tensor\n",
    "        elif operation == '+':\n",
    "            return first_tensor + second_tensor\n",
    "        elif operation == '-':\n",
    "            return first_tensor - second_tensor\n",
    "        else:\n",
    "            raise ConfigurationError(\"Invalid operation: \" + operation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_combination_dim(combination: str, tensor_dims) -> int:\n",
    "    if combination.isdigit():\n",
    "        index = int(combination) - 1\n",
    "        return tensor_dims[index]\n",
    "    else:\n",
    "        if len(combination) != 3:\n",
    "            raise ConfigurationError(\"Invalid combination: \" + combination)\n",
    "        first_tensor_dim = _get_combination_dim(combination[0], tensor_dims)\n",
    "        second_tensor_dim = _get_combination_dim(combination[2], tensor_dims)\n",
    "        operation = combination[1]\n",
    "        if first_tensor_dim != second_tensor_dim:\n",
    "            raise ConfigurationError(\"Tensor dims must match for operation \\\"{}\\\"\".format(operation))\n",
    "        return first_tensor_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weighted_sum(matrix: torch.Tensor, attention: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Takes a matrix of vectors and a set of weights over the rows in the matrix (which we call an\n",
    "    \"attention\" vector), and returns a weighted sum of the rows in the matrix.  This is the typical\n",
    "    computation performed after an attention mechanism.\n",
    "    Note that while we call this a \"matrix\" of vectors and an attention \"vector\", we also handle\n",
    "    higher-order tensors.  We always sum over the second-to-last dimension of the \"matrix\", and we\n",
    "    assume that all dimensions in the \"matrix\" prior to the last dimension are matched in the\n",
    "    \"vector\".  Non-matched dimensions in the \"vector\" must be `directly after the batch dimension`.\n",
    "    For example, say I have a \"matrix\" with dimensions ``(batch_size, num_queries, num_words,\n",
    "    embedding_dim)``.  The attention \"vector\" then must have at least those dimensions, and could\n",
    "    have more. Both:\n",
    "        - ``(batch_size, num_queries, num_words)`` (distribution over words for each query)\n",
    "        - ``(batch_size, num_documents, num_queries, num_words)`` (distribution over words in a\n",
    "          query for each document)\n",
    "    are valid input \"vectors\", producing tensors of shape:\n",
    "    ``(batch_size, num_queries, embedding_dim)`` and\n",
    "    ``(batch_size, num_documents, num_queries, embedding_dim)`` respectively.\n",
    "    \"\"\"\n",
    "    # We'll special-case a few settings here, where there are efficient (but poorly-named)\n",
    "    # operations in pytorch that already do the computation we need.\n",
    "    if attention.dim() == 2 and matrix.dim() == 3:\n",
    "        return attention.unsqueeze(1).bmm(matrix).squeeze(1)\n",
    "    if attention.dim() == 3 and matrix.dim() == 3:\n",
    "        return attention.bmm(matrix)\n",
    "    if matrix.dim() - 1 < attention.dim():\n",
    "        expanded_size = list(matrix.size())\n",
    "        for i in range(attention.dim() - matrix.dim() + 1):\n",
    "            matrix = matrix.unsqueeze(1)\n",
    "            expanded_size.insert(i + 1, attention.size(i + 1))\n",
    "        matrix = matrix.expand(*expanded_size)\n",
    "    intermediate = attention.unsqueeze(-1).expand_as(matrix) * matrix\n",
    "    return intermediate.sum(dim=-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TimeDistributed(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Given an input shaped like ``(batch_size, time_steps, [rest])`` and a ``Module`` that takes\n",
    "    inputs like ``(batch_size, [rest])``, ``TimeDistributed`` reshapes the input to be\n",
    "    ``(batch_size * time_steps, [rest])``, applies the contained ``Module``, then reshapes it back.\n",
    "    Note that while the above gives shapes with ``batch_size`` first, this ``Module`` also works if\n",
    "    ``batch_size`` is second - we always just combine the first two dimensions, then split them.\n",
    "    \"\"\"\n",
    "    def __init__(self, module):\n",
    "        super(TimeDistributed, self).__init__()\n",
    "        self._module = module\n",
    "\n",
    "    def forward(self, *inputs):  # pylint: disable=arguments-differ\n",
    "        reshaped_inputs = []\n",
    "        for input_tensor in inputs:\n",
    "            input_size = input_tensor.size()\n",
    "            if len(input_size) <= 2:\n",
    "                raise RuntimeError(\"No dimension to distribute: \" + str(input_size))\n",
    "\n",
    "            # Squash batch_size and time_steps into a single axis; result has shape\n",
    "            # (batch_size * time_steps, input_size).\n",
    "            squashed_shape = [-1] + [x for x in input_size[2:]]\n",
    "            reshaped_inputs.append(input_tensor.contiguous().view(*squashed_shape))\n",
    "\n",
    "        reshaped_outputs = self._module(*reshaped_inputs)\n",
    "\n",
    "        # Now get the output back into the right shape.\n",
    "        # (batch_size, time_steps, [hidden_size])\n",
    "        new_shape = [input_size[0], input_size[1]] + [x for x in reshaped_outputs.size()[1:]]\n",
    "        outputs = reshaped_outputs.contiguous().view(*new_shape)\n",
    "\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_best_span(span_start_logits: Variable, span_end_logits: Variable) -> Variable:\n",
    "        if span_start_logits.dim() != 2 or span_end_logits.dim() != 2:\n",
    "            raise ValueError(\"Input shapes must be (batch_size, passage_length)\")\n",
    "        batch_size, passage_length = span_start_logits.size()\n",
    "        max_span_log_prob = [-1e20] * batch_size\n",
    "        span_start_argmax = [0] * batch_size\n",
    "        best_word_span = Variable(span_start_logits.data.new()\n",
    "                                  .resize_(batch_size, 2).fill_(0)).long()\n",
    "\n",
    "        span_start_logits = span_start_logits.data.cpu().numpy()\n",
    "        span_end_logits = span_end_logits.data.cpu().numpy()\n",
    "\n",
    "        for b in range(batch_size):  # pylint: disable=invalid-name\n",
    "            for j in range(passage_length):\n",
    "                val1 = span_start_logits[b, span_start_argmax[b]]\n",
    "                if val1 < span_start_logits[b, j]:\n",
    "                    span_start_argmax[b] = j\n",
    "                    val1 = span_start_logits[b, j]\n",
    "\n",
    "                val2 = span_end_logits[b, j]\n",
    "\n",
    "                if val1 + val2 > max_span_log_prob[b]:\n",
    "                    best_word_span[b, 0] = span_start_argmax[b]\n",
    "                    best_word_span[b, 1] = j\n",
    "                    max_span_log_prob[b] = val1 + val2\n",
    "        return best_word_span"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"\\nTODO THIS ONLY WORKS 1 WAY, EITHER MAKE BIDIRECTIONAL ORRRRRRR HAVE AN INPUT WITH THE WORDS REVERSED\\nI DON'T ACTUALLY THINK YOU NEED TO MAKE BIDIRECTIONAL NOW?\\n\""
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = Variable(torch.LongTensor(np.random.randint(0,50,size=(32))))\n",
    "#50 is vocab size, 64 is embedding size\n",
    "\"\"\"\n",
    "TODO THIS ONLY WORKS 1 WAY, EITHER MAKE BIDIRECTIONAL ORRRRRRR HAVE AN INPUT WITH THE WORDS REVERSED\n",
    "I DON'T ACTUALLY THINK YOU NEED TO MAKE BIDIRECTIONAL NOW?\n",
    "\"\"\"\n",
    "\n",
    "#highway = Highway(50, 200, 2)\n",
    "#highway(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_embedding_dim = 100\n",
    "char_embedding_dim = 16\n",
    "n_words = 1000 #vocab size, not length\n",
    "n_chars = 256 #vocab size, not length\n",
    "batch_size = 32\n",
    "context_max_word_len = 250 #max words in context\n",
    "context_max_char_len = 10 #max characters per word\n",
    "query_max_word_len = 15 #max words in context\n",
    "query_max_char_len = 10 #max characters per word SHOULD BE THE SAME AS CONTEXT_MAX_CHAR_LEN\n",
    "\n",
    "context_words = Variable(torch.LongTensor(np.random.randint(0,n_words,size=(batch_size,context_max_word_len))))\n",
    "context_chars = Variable(torch.LongTensor(np.random.randint(0,n_chars,size=(batch_size,context_max_word_len,context_max_char_len))))\n",
    "\n",
    "query_words = Variable(torch.LongTensor(np.random.randint(0,n_words,size=(batch_size,query_max_word_len))))\n",
    "query_chars = Variable(torch.LongTensor(np.random.randint(0,n_chars,size=(batch_size,query_max_word_len,query_max_char_len))))\n",
    "\n",
    "\"\"\"\n",
    "BEGIN CONTEXT TO ATTENTION FLOW LAYER INPUTS\n",
    "x_T to h_T\n",
    "\"\"\"\n",
    "\n",
    "word_embedding = Embedding(n_words, word_embedding_dim) #instantiate word -> word vectors module\n",
    "char_embedding = TimeDistributed(Embedding(n_chars, char_embedding_dim)) #instantiate char -> word vectors module\n",
    "\n",
    "c_embedded_word = word_embedding(context_words) #input words, get out vectors\n",
    "c_embedded_char = char_embedding(context_chars) #input chars, get out vectors\n",
    "\n",
    "char_cnn = TimeDistributed(CNN(n_chars, char_embedding_dim)) #instantiate word vectors -> cnn vectors module\n",
    "\n",
    "c_cnn_embedded_chars = char_cnn(c_embedded_char) #input chars, get out vectors\n",
    "\n",
    "highway = TimeDistributed(Highway(word_embedding_dim*2, num_layers=2)) #instantiate word_emb + char_emb (from CNN) -> embedded\n",
    "\n",
    "c_highway_input = torch.cat((c_cnn_embedded_chars, c_embedded_word), dim=2) #concat char_cnn_emb + word_emb\n",
    "\n",
    "c_embedded = highway(c_highway_input) #input concat, get out embedded context\n",
    "\n",
    "phrase_layer = LSTM(word_embedding_dim*2, word_embedding_dim, n_layers=1, bidirectional=True) #lstm to turn context embedding into phrase embedding\n",
    "\n",
    "c_embedded_phrase = phrase_layer(c_embedded) #pass through layer to get embedded phrase\n",
    "\n",
    "\"\"\"\n",
    "BEGIN QUERY TO ATTENTION FLOW LAYER INPUTS\n",
    "q_J to u_J\n",
    "\"\"\"\n",
    "\n",
    "q_embedded_word = word_embedding(query_words) #input words, get out vectors\n",
    "q_embedded_char = char_embedding(query_chars) #input chars, get out vectors\n",
    "\n",
    "q_cnn_embedded_chars = char_cnn(q_embedded_char) #input chars, get out vectors\n",
    "\n",
    "q_highway_input = torch.cat((q_cnn_embedded_chars, q_embedded_word), dim=2) ##concat char_cnn_emb + word_emb\n",
    "\n",
    "q_embedded = highway(q_highway_input) #input concat, get out embedded context\n",
    "\n",
    "q_embedded_phrase = phrase_layer(q_embedded) #pass through layer to get embedded phrase\n",
    "\n",
    "similarity_function = LinearSimilarity(200,200) #used for matrix attention\n",
    "\n",
    "attention_matrix = MatrixAttention(similarity_function) #gets un-normalized matrix attention\n",
    "\n",
    "#must be [batch, length, emb]\n",
    "c_embedded_phrase = c_embedded_phrase.permute(1, 0, 2)\n",
    "q_embedded_phrase = q_embedded_phrase.permute(1, 0, 2)\n",
    "\n",
    "\"\"\"\n",
    "ATTENTION FLOW LAYER\n",
    "\"\"\"\n",
    "\n",
    "c_q_similarity = attention_matrix(c_embedded_phrase, q_embedded_phrase) #similarity between context and query\n",
    "\n",
    "c_q_attention = F.softmax(c_q_similarity, dim=2) #normalise along the query dim (the shorter one) \n",
    "\n",
    "c_q_vectors = weighted_sum(q_embedded_phrase, c_q_attention) #apply attention to query\n",
    "\n",
    "q_c_similarity = c_q_similarity.max(dim=-1)[0]\n",
    "\n",
    "q_c_attention = F.softmax(q_c_similarity, dim=1) #normalise along the context dim (the longer one)\n",
    "\n",
    "q_c_vectors = weighted_sum(c_embedded_phrase, q_c_attention) #apply attention to context\n",
    "\n",
    "tiled_q_c_vectors = q_c_vectors.unsqueeze(1).expand(c_q_vectors.shape) #make weighted vectors same size\n",
    "\n",
    "#have: \n",
    "#the emb phrase context\n",
    "#context attn applied to query\n",
    "#emb phrase context * context attn applied to query \n",
    "#emb phrase context * query attn applied to context\n",
    "final_merged_c = torch.cat([c_embedded_phrase,\n",
    "                            c_q_vectors,\n",
    "                            c_embedded_phrase * c_q_vectors,\n",
    "                            c_embedded_phrase * tiled_q_c_vectors],\n",
    "                            dim=-1)\n",
    "\n",
    "modeling_layer = LSTM(word_embedding_dim*8, word_embedding_dim, n_layers=2, bidirectional=True) #uses attention and phrase to \"model\" \n",
    "\n",
    "modeled_context = modeling_layer(final_merged_c) #apply modeling layer\n",
    "\n",
    "modeled_context = modeled_context.permute(1, 0, 2) #need to permute after lstm for batch first\n",
    "\n",
    "span_start_input = torch.cat((final_merged_c, modeled_context), dim=-1) #embedded phrase and attention concat modeling\n",
    "\n",
    "span_start_predictor = TimeDistributed(nn.Linear(1000, 1)) #module to predict span start from phrase + attention\n",
    "\n",
    "span_start_logits = span_start_predictor(span_start_input).squeeze(-1) #apply module to predict span start from phrase + attention\n",
    "\n",
    "span_start_probs = F.softmax(span_start_logits, dim=1) #turn into probabilities over context to where answer span begins\n",
    "\n",
    "span_start_representation = weighted_sum(modeled_context, span_start_probs) #apply the probabilities over the modeled context\n",
    "\n",
    "tiled_start_representation = span_start_representation.unsqueeze(1).expand(modeled_context.shape) #reshape to predict span end\n",
    "\n",
    "#have:\n",
    "#phrase + attention merge\n",
    "#modeled context from the above\n",
    "#begin of span representation\n",
    "#above 2 mulitplied together\n",
    "span_end_representation = torch.cat([final_merged_c,\n",
    "                                     modeled_context,\n",
    "                                     tiled_start_representation,\n",
    "                                     modeled_context * tiled_start_representation],\n",
    "                                     dim=-1)\n",
    "\n",
    "#LSTM over all of the above\n",
    "span_end_encoder = LSTM(1400, 100, n_layers=1, bidirectional=True)\n",
    "\n",
    "#used to find span end\n",
    "encoded_span_end = span_end_encoder(span_end_representation)\n",
    "\n",
    "encoded_span_end = encoded_span_end.permute(1, 0, 2)\n",
    "\n",
    "#conat lstm representation as well as all the phrase and context stuff\n",
    "span_end_input = torch.cat((final_merged_c, encoded_span_end),dim=-1)\n",
    "\n",
    "#module to produce logits \n",
    "span_end_predictor = TimeDistributed(torch.nn.Linear(1000, 1))\n",
    "\n",
    "#un-normalized outputs for span end probabily\n",
    "span_end_logits = span_end_predictor(span_end_input).squeeze(-1)\n",
    "\n",
    "#prob of each word in span being the span end\n",
    "span_end_probs = F.softmax(span_end_logits, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable containing:\n",
      "1.00000e-03 *\n",
      " 3.9767  4.0361  4.0894  ...   4.0464  3.8924  3.9780\n",
      " 3.8900  4.0661  4.0770  ...   4.0577  4.0544  3.8594\n",
      " 3.9607  3.9651  3.8803  ...   3.8963  4.0381  4.0798\n",
      "          ...             ⋱             ...          \n",
      " 3.9333  4.0308  3.9659  ...   4.1081  4.0074  3.9999\n",
      " 3.7908  3.9123  4.1663  ...   3.8650  3.9431  4.0505\n",
      " 3.9692  3.9894  4.0286  ...   4.0498  4.0669  4.1005\n",
      "[torch.FloatTensor of size 32x250]\n",
      "\n",
      "Variable containing:\n",
      "1.00000e-03 *\n",
      " 3.8633  3.9326  3.9281  ...   3.9944  4.0016  3.8890\n",
      " 3.9752  3.9891  4.1079  ...   4.0001  3.9269  3.9458\n",
      " 3.9494  3.9447  4.0215  ...   4.0203  4.1238  4.0914\n",
      "          ...             ⋱             ...          \n",
      " 4.0060  4.0729  4.0784  ...   4.0225  3.9789  3.9763\n",
      " 3.9401  3.8875  4.1226  ...   3.9323  3.9203  4.0247\n",
      " 3.9129  3.9043  3.8974  ...   3.8401  3.8216  3.9164\n",
      "[torch.FloatTensor of size 32x250]\n",
      "\n",
      "Variable containing:\n",
      "   45    73\n",
      "   59   194\n",
      "   45    45\n",
      "   93    93\n",
      "   95   173\n",
      "  100   139\n",
      "  145   241\n",
      "   43   170\n",
      "  116   117\n",
      "   93   245\n",
      "   18   221\n",
      "   28    84\n",
      "    6    36\n",
      "   19   160\n",
      "   31   141\n",
      "  135   225\n",
      "  125   175\n",
      "  129   141\n",
      "   84   245\n",
      "  155   167\n",
      "  184   221\n",
      "   20    35\n",
      "   61   129\n",
      "    5    28\n",
      "  185   185\n",
      "   68   161\n",
      "   70   125\n",
      "  162   234\n",
      "  147   149\n",
      "  131   239\n",
      "  173   203\n",
      "  145   207\n",
      "[torch.LongTensor of size 32x2]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(span_start_probs)\n",
    "print(span_end_probs)\n",
    "print(_get_best_span(span_start_probs, span_end_probs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}