{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.layers.attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attention Layers\n",
    "> Implementation of Attention modules including Multihead attention etc.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TokenEmbedding(nn.Embedding):\n",
    "    def __init__(self, vocab_size, embed_size=512):\n",
    "        super().__init__(vocab_size, embed_size, padding_idx=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class PositionalEmbedding(nn.Module):\n",
    "    def __init__(self, max_len, d_model):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.pe = nn.Embedding(max_len+1, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        pose = (x > 0) * (x > 0).sum(dim=-1).unsqueeze(1).repeat(1, x.size(-1))\n",
    "        pose += torch.arange(start=-(x.size(1)-1), end=1, step=1, device=x.device)\n",
    "        pose = pose * (x > 0)\n",
    "\n",
    "        return self.pe(pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class GELU(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class PositionwiseFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, d_ff):\n",
    "        super().__init__()\n",
    "        self.w_1 = nn.Linear(d_model, d_ff)\n",
    "        self.w_2 = nn.Linear(d_ff, d_model)\n",
    "        self.activation = GELU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.w_2(self.activation(self.w_1(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, features, eps=1e-6):\n",
    "        super().__init__()\n",
    "        self.weight = nn.Parameter(torch.ones(features))\n",
    "        self.bias = nn.Parameter(torch.zeros(features))\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        return self.weight * (x - mean) / (std + self.eps) + self.bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SublayerConnection(nn.Module):\n",
    "    \"\"\"layer norm and dropout (dropout and then layer norm)\n",
    "    \"\"\"\n",
    "    def __init__(self, size, dropout):\n",
    "        super().__init__()\n",
    "        self.layer_norm = LayerNorm(size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, sublayer):\n",
    "        # return x + self.dropout(sublayer(self.norm(x)))  # original implementation\n",
    "        return self.layer_norm(x + self.dropout(sublayer(x)))  # BERT4Rec implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class Attention(nn.Module):\n",
    "    def forward(self, query, key, value, mask=None, dropout=None, sas=False):\n",
    "        scores = torch.matmul(query, key.transpose(-2, -1)) \\\n",
    "            / math.sqrt(query.size(-1))\n",
    "\n",
    "        if mask is not None:\n",
    "            scores = scores.masked_fill(mask == 0, -1e9)\n",
    "\n",
    "        if sas:\n",
    "            direction_mask = torch.ones_like(scores)\n",
    "            direction_mask = torch.tril(direction_mask)\n",
    "            scores = scores.masked_fill(direction_mask == 0, -1e9)\n",
    "\n",
    "        p_attn = F.softmax(scores, dim=-1)\n",
    "\n",
    "        if dropout is not None:\n",
    "            p_attn = dropout(p_attn)\n",
    "\n",
    "        return torch.matmul(p_attn, value), p_attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class MultiHeadedAttention(nn.Module):\n",
    "    def __init__(self, h, d_model, head_size=None, dropout=0.1):\n",
    "        super().__init__()\n",
    "        assert d_model % h == 0\n",
    "\n",
    "        self.h = h\n",
    "        self.d_k = d_model // h\n",
    "        if head_size is not None:\n",
    "            self.head_size = head_size\n",
    "        else:\n",
    "            self.head_size = d_model // h\n",
    "\n",
    "        self.linear_layers = nn.ModuleList(\n",
    "            [nn.Linear(d_model, self.h * self.head_size) for _ in range(3)])\n",
    "        self.attention = Attention()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        self.output_linear = nn.Linear(self.h * self.head_size, d_model)\n",
    "\n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        batch_size = query.size(0)\n",
    "\n",
    "        # 1) do all the linear projections in batch from d_model => h x d_k\n",
    "        query, key, value = [l(x).view(batch_size, -1, self.h, self.head_size).transpose(1, 2)\n",
    "                             for l, x in zip(self.linear_layers, (query, key, value))]\n",
    "        \n",
    "        # 2) apply attention on all the projected vectors in batch.\n",
    "        x, attn = self.attention(\n",
    "            query, key, value, mask=mask, dropout=self.dropout)\n",
    "\n",
    "        # 3) \"concat\" using a view and apply a final linear.\n",
    "        x = x.transpose(1, 2).contiguous().view(\n",
    "            batch_size, -1, self.h * self.head_size)\n",
    "        return self.output_linear(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, hidden, attn_heads, head_size, feed_forward_hidden, dropout, attn_dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.attention = MultiHeadedAttention(\n",
    "            h=attn_heads, d_model=hidden, head_size=head_size, dropout=attn_dropout)\n",
    "        self.feed_forward = PositionwiseFeedForward(\n",
    "            d_model=hidden, d_ff=feed_forward_hidden)\n",
    "        self.input_sublayer = SublayerConnection(size=hidden, dropout=dropout)\n",
    "        self.output_sublayer = SublayerConnection(size=hidden, dropout=dropout)\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        x = self.input_sublayer(\n",
    "            x, lambda _x: self.attention.forward(_x, _x, _x, mask=mask))\n",
    "        x = self.output_sublayer(x, self.feed_forward)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SASMultiHeadedAttention(nn.Module):\n",
    "    def __init__(self, h, d_model, head_size=None, dropout=0.1):\n",
    "        super().__init__()\n",
    "        assert d_model % h == 0\n",
    "\n",
    "        self.h = h\n",
    "        self.d_k = d_model // h\n",
    "        if head_size is not None:\n",
    "            self.head_size = head_size\n",
    "        else:\n",
    "            self.head_size = d_model // h\n",
    "\n",
    "        self.linear_layers = nn.ModuleList(\n",
    "            [nn.Linear(d_model, self.h * self.head_size) for _ in range(3)])\n",
    "        self.attention = Attention()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        self.layer_norm = LayerNorm(d_model)\n",
    "\n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        batch_size = query.size(0)\n",
    "\n",
    "        # 1) do all the linear projections in batch from d_model => h x d_k\n",
    "        query_, key_, value_ = [l(x).view(batch_size, -1, self.h, self.head_size).transpose(1, 2)\n",
    "                             for l, x in zip(self.linear_layers, (query, key, value))]\n",
    "        \n",
    "        # 2) apply attention on all the projected vectors in batch.\n",
    "        x, attn = self.attention(\n",
    "            query_, key_, value_, mask=mask, dropout=self.dropout, sas=True)\n",
    "\n",
    "        # 3) \"concat\" using a view and apply a final linear.\n",
    "        x = x.transpose(1, 2).contiguous().view(\n",
    "            batch_size, -1, self.h * self.head_size)\n",
    "        \n",
    "        return self.layer_norm(x + query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SASPositionwiseFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, d_ff, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)\n",
    "        self.activation = nn.ReLU()\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)\n",
    "        self.layer_norm = LayerNorm(d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_ = self.dropout(self.activation(self.conv1(x.permute(0, 2, 1))))\n",
    "        return self.layer_norm(self.dropout(self.conv2(x_)).permute(0, 2, 1) + x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SASTransformerBlock(nn.Module):\n",
    "    def __init__(self, hidden, attn_heads, head_size, feed_forward_hidden, dropout, attn_dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.layer_norm = LayerNorm(hidden)\n",
    "        self.attention = SASMultiHeadedAttention(\n",
    "            h=attn_heads, d_model=hidden, head_size=head_size, dropout=attn_dropout)\n",
    "        self.feed_forward = SASPositionwiseFeedForward(\n",
    "            d_model=hidden, d_ff=feed_forward_hidden, dropout=dropout)\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        x = self.attention(self.layer_norm(x), x, x, mask)\n",
    "        x = self.feed_forward(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> References\n",
    "1. https://github.com/sparsh-ai/stanza/blob/S714864/model/attention.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-30 08:23:40\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.144+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "IPython: 5.5.0\n",
      "torch  : 1.10.0+cu111\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
