{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.bert"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT\n",
    "> Bidirectional Encoder Representations from Transformers.\n",
    "\n",
    "BERT was one of the first autoencoding language models to utilize the encoder Transformer stack with slight modifications for language modeling. Its architecture is a multilayer Transformer encoder based on the Transformer original implementation. The Transformer model itself was originally for machine translation tasks, but the main improvement made by BERT is the utilization of this part of the architecture to provide better language modeling. This language model, after pretraining, is able to provide a global understanding of the language it is trained on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.nb_imports import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "from torch import nn as nn\n",
    "import math\n",
    "\n",
    "from recohut.models.layers.attention import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class BERTEmbedding(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        vocab_size = args.num_items + 2\n",
    "        hidden = args.bert_hidden_units\n",
    "        max_len = args.bert_max_len\n",
    "        dropout = args.bert_dropout\n",
    "\n",
    "        self.token = TokenEmbedding(\n",
    "            vocab_size=vocab_size, embed_size=hidden)\n",
    "        self.position = PositionalEmbedding(\n",
    "            max_len=max_len, d_model=hidden)\n",
    "\n",
    "        self.layer_norm = LayerNorm(features=hidden)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def get_mask(self, x):\n",
    "        if len(x.shape) > 2:\n",
    "            x = torch.ones(x.shape[:2]).to(x.device)\n",
    "        return (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        mask = self.get_mask(x)\n",
    "        if len(x.shape) > 2:\n",
    "            pos = self.position(torch.ones(x.shape[:2]).to(x.device))\n",
    "            x = torch.matmul(x, self.token.weight) + pos\n",
    "        else:\n",
    "            x = self.token(x) + self.position(x)\n",
    "        return self.dropout(self.layer_norm(x)), mask\n",
    "\n",
    "\n",
    "class BERTModel(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        hidden = args.bert_hidden_units\n",
    "        heads = args.bert_num_heads\n",
    "        head_size = args.bert_head_size\n",
    "        dropout = args.bert_dropout\n",
    "        attn_dropout = args.bert_attn_dropout\n",
    "        layers = args.bert_num_blocks\n",
    "\n",
    "        self.transformer_blocks = nn.ModuleList([TransformerBlock(\n",
    "            hidden, heads, head_size, hidden * 4, dropout, attn_dropout) for _ in range(layers)])\n",
    "        self.linear = nn.Linear(hidden, hidden)\n",
    "        self.bias = torch.nn.Parameter(torch.zeros(args.num_items + 2))\n",
    "        self.bias.requires_grad = True\n",
    "        self.activation = GELU()\n",
    "\n",
    "    def forward(self, x, embedding_weight, mask):\n",
    "        for transformer in self.transformer_blocks:\n",
    "            x = transformer.forward(x, mask)\n",
    "        x = self.activation(self.linear(x))\n",
    "        scores = torch.matmul(x, embedding_weight.permute(1, 0)) + self.bias\n",
    "        return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class BERT(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        self.args = args\n",
    "        self.embedding = BERTEmbedding(self.args)\n",
    "        self.model = BERTModel(self.args)\n",
    "        self.truncated_normal_init()\n",
    "\n",
    "    def truncated_normal_init(self, mean=0, std=0.02, lower=-0.04, upper=0.04):\n",
    "        with torch.no_grad():\n",
    "            l = (1. + math.erf(((lower - mean) / std) / math.sqrt(2.))) / 2.\n",
    "            u = (1. + math.erf(((upper - mean) / std) / math.sqrt(2.))) / 2.\n",
    "\n",
    "            for n, p in self.model.named_parameters():\n",
    "                if not 'layer_norm' in n:\n",
    "                    p.uniform_(2 * l - 1, 2 * u - 1)\n",
    "                    p.erfinv_()\n",
    "                    p.mul_(std * math.sqrt(2.))\n",
    "                    p.add_(mean)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x, mask = self.embedding(x)\n",
    "        scores = self.model(x, self.embedding.token.weight, mask)\n",
    "        return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<bound method Module.parameters of BERT(\n",
       "  (embedding): BERTEmbedding(\n",
       "    (token): TokenEmbedding(12, 4, padding_idx=0)\n",
       "    (position): PositionalEmbedding(\n",
       "      (pe): Embedding(9, 4)\n",
       "    )\n",
       "    (layer_norm): LayerNorm()\n",
       "    (dropout): Dropout(p=0.2, inplace=False)\n",
       "  )\n",
       "  (model): BERTModel(\n",
       "    (transformer_blocks): ModuleList(\n",
       "      (0): TransformerBlock(\n",
       "        (attention): MultiHeadedAttention(\n",
       "          (linear_layers): ModuleList(\n",
       "            (0): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (1): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (2): Linear(in_features=4, out_features=8, bias=True)\n",
       "          )\n",
       "          (attention): Attention()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (output_linear): Linear(in_features=8, out_features=4, bias=True)\n",
       "        )\n",
       "        (feed_forward): PositionwiseFeedForward(\n",
       "          (w_1): Linear(in_features=4, out_features=16, bias=True)\n",
       "          (w_2): Linear(in_features=16, out_features=4, bias=True)\n",
       "          (activation): GELU()\n",
       "        )\n",
       "        (input_sublayer): SublayerConnection(\n",
       "          (layer_norm): LayerNorm()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "        )\n",
       "        (output_sublayer): SublayerConnection(\n",
       "          (layer_norm): LayerNorm()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (1): TransformerBlock(\n",
       "        (attention): MultiHeadedAttention(\n",
       "          (linear_layers): ModuleList(\n",
       "            (0): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (1): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (2): Linear(in_features=4, out_features=8, bias=True)\n",
       "          )\n",
       "          (attention): Attention()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (output_linear): Linear(in_features=8, out_features=4, bias=True)\n",
       "        )\n",
       "        (feed_forward): PositionwiseFeedForward(\n",
       "          (w_1): Linear(in_features=4, out_features=16, bias=True)\n",
       "          (w_2): Linear(in_features=16, out_features=4, bias=True)\n",
       "          (activation): GELU()\n",
       "        )\n",
       "        (input_sublayer): SublayerConnection(\n",
       "          (layer_norm): LayerNorm()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "        )\n",
       "        (output_sublayer): SublayerConnection(\n",
       "          (layer_norm): LayerNorm()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (2): TransformerBlock(\n",
       "        (attention): MultiHeadedAttention(\n",
       "          (linear_layers): ModuleList(\n",
       "            (0): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (1): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (2): Linear(in_features=4, out_features=8, bias=True)\n",
       "          )\n",
       "          (attention): Attention()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (output_linear): Linear(in_features=8, out_features=4, bias=True)\n",
       "        )\n",
       "        (feed_forward): PositionwiseFeedForward(\n",
       "          (w_1): Linear(in_features=4, out_features=16, bias=True)\n",
       "          (w_2): Linear(in_features=16, out_features=4, bias=True)\n",
       "          (activation): GELU()\n",
       "        )\n",
       "        (input_sublayer): SublayerConnection(\n",
       "          (layer_norm): LayerNorm()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "        )\n",
       "        (output_sublayer): SublayerConnection(\n",
       "          (layer_norm): LayerNorm()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (3): TransformerBlock(\n",
       "        (attention): MultiHeadedAttention(\n",
       "          (linear_layers): ModuleList(\n",
       "            (0): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (1): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (2): Linear(in_features=4, out_features=8, bias=True)\n",
       "          )\n",
       "          (attention): Attention()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (output_linear): Linear(in_features=8, out_features=4, bias=True)\n",
       "        )\n",
       "        (feed_forward): PositionwiseFeedForward(\n",
       "          (w_1): Linear(in_features=4, out_features=16, bias=True)\n",
       "          (w_2): Linear(in_features=16, out_features=4, bias=True)\n",
       "          (activation): GELU()\n",
       "        )\n",
       "        (input_sublayer): SublayerConnection(\n",
       "          (layer_norm): LayerNorm()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "        )\n",
       "        (output_sublayer): SublayerConnection(\n",
       "          (layer_norm): LayerNorm()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (linear): Linear(in_features=4, out_features=4, bias=True)\n",
       "    (activation): GELU()\n",
       "  )\n",
       ")>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Args:\n",
    "    bert_hidden_units = 4\n",
    "    bert_num_heads = 2\n",
    "    bert_head_size = 4\n",
    "    bert_dropout = 0.2\n",
    "    bert_attn_dropout = 0.2\n",
    "    bert_num_blocks = 4\n",
    "    num_items = 10\n",
    "    bert_hidden_units = 4\n",
    "    bert_max_len = 8\n",
    "    bert_dropout = 0.2\n",
    "\n",
    "args = Args()\n",
    "model = BERT(args)\n",
    "model.parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> References\n",
    "1. https://github.com/Yueeeeeeee/RecSys-Extraction-Attack/blob/main/model/bert.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-31 07:01:15\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",
      "numpy     : 1.19.5\n",
      "matplotlib: 3.2.2\n",
      "torch     : 1.10.0+cu111\n",
      "PIL       : 7.1.2\n",
      "IPython   : 5.5.0\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
}
