BERT-MLM/masked-language-modeling.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f1b35fb6",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:29.185259Z",
     "iopub.status.busy": "2023-10-09T16:24:29.183945Z",
     "iopub.status.idle": "2023-10-09T16:24:32.823094Z",
     "shell.execute_reply": "2023-10-09T16:24:32.821964Z"
    },
    "papermill": {
     "duration": 3.652306,
     "end_time": "2023-10-09T16:24:32.826006",
     "exception": false,
     "start_time": "2023-10-09T16:24:29.173700",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "d6486b6d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:32.840432Z",
     "iopub.status.busy": "2023-10-09T16:24:32.839978Z",
     "iopub.status.idle": "2023-10-09T16:24:32.846881Z",
     "shell.execute_reply": "2023-10-09T16:24:32.845446Z"
    },
    "papermill": {
     "duration": 0.016688,
     "end_time": "2023-10-09T16:24:32.849188",
     "exception": false,
     "start_time": "2023-10-09T16:24:32.832500",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def sample_input_ids():\n",
    "    input_ids = torch.randint(0,30_000,(4,20))\n",
    "    input_ids[:,0]=101 # CLS TOKEN\n",
    "    input_ids[:,7]=103 # SEP TOKEN\n",
    "    input_ids[:,-1]=103 # SEP TOKEN\n",
    "    \n",
    "    input_ids[:,16:]=0 # PAD TOKEN\n",
    "    return input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "a5427cc6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:32.862639Z",
     "iopub.status.busy": "2023-10-09T16:24:32.862200Z",
     "iopub.status.idle": "2023-10-09T16:24:32.940985Z",
     "shell.execute_reply": "2023-10-09T16:24:32.939433Z"
    },
    "papermill": {
     "duration": 0.089655,
     "end_time": "2023-10-09T16:24:32.944539",
     "exception": false,
     "start_time": "2023-10-09T16:24:32.854884",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  101,  3413,  8401,  4797, 13763, 22724,  7058,   103, 24223,  9199,\n",
       "         11515,   488, 27828,  3547, 18158, 10189,     0,     0,     0,     0],\n",
       "        [  101,  4303, 15277,  4807, 26430, 21759, 24603,   103, 12689, 15363,\n",
       "         14318,     5,  9876, 29946,  4058,   451,     0,     0,     0,     0],\n",
       "        [  101, 12024, 27899,  8706, 19037, 28294, 16359,   103,   835, 13595,\n",
       "         20214, 23694,  8276, 11985, 14855,  8413,     0,     0,     0,     0],\n",
       "        [  101, 15900,   527, 15302, 26235, 16148, 10709,   103, 21784, 18223,\n",
       "         13851, 11017, 22713, 25382, 20000, 24657,     0,     0,     0,     0]])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids = sample_input_ids()\n",
    "input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "403b3bcf",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:32.960603Z",
     "iopub.status.busy": "2023-10-09T16:24:32.960187Z",
     "iopub.status.idle": "2023-10-09T16:24:32.991066Z",
     "shell.execute_reply": "2023-10-09T16:24:32.990317Z"
    },
    "papermill": {
     "duration": 0.042049,
     "end_time": "2023-10-09T16:24:32.993440",
     "exception": false,
     "start_time": "2023-10-09T16:24:32.951391",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[False, False, False, False,  True, False, False, False, False, False,\n",
       "         False, False, False, False, False, False, False, False, False, False],\n",
       "        [False, False, False,  True, False, False, False, False, False, False,\n",
       "         False, False, False,  True, False, False, False, False, False, False],\n",
       "        [False, False, False, False,  True, False, False, False, False, False,\n",
       "         False,  True, False, False, False, False, False, False, False, False],\n",
       "        [False, False, False, False, False, False, False, False, False, False,\n",
       "         False, False,  True, False,  True,  True, False, False, False, False]])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mlm_mask = torch.rand(input_ids.size()) < 0.15 * (input_ids!=101) * (input_ids!=103) * (input_ids!=0) \n",
    "mlm_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "26969e33",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.014338Z",
     "iopub.status.busy": "2023-10-09T16:24:33.013440Z",
     "iopub.status.idle": "2023-10-09T16:24:33.023948Z",
     "shell.execute_reply": "2023-10-09T16:24:33.022782Z"
    },
    "papermill": {
     "duration": 0.026718,
     "end_time": "2023-10-09T16:24:33.026553",
     "exception": false,
     "start_time": "2023-10-09T16:24:32.999835",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[    0,     0,     0,     0, 13763,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0],\n",
       "        [    0,     0,     0,  4807,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0, 29946,     0,     0,     0,     0,     0,     0],\n",
       "        [    0,     0,     0,     0, 19037,     0,     0,     0,     0,     0,\n",
       "             0, 23694,     0,     0,     0,     0,     0,     0,     0,     0],\n",
       "        [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0, 22713,     0, 20000, 24657,     0,     0,     0,     0]])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "masked_tokens = input_ids * mlm_mask\n",
    "masked_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "dba1f51e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.041415Z",
     "iopub.status.busy": "2023-10-09T16:24:33.041051Z",
     "iopub.status.idle": "2023-10-09T16:24:33.054776Z",
     "shell.execute_reply": "2023-10-09T16:24:33.053586Z"
    },
    "papermill": {
     "duration": 0.025121,
     "end_time": "2023-10-09T16:24:33.058184",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.033063",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_ids[masked_tokens!=0]=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "4b29ba5c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.073989Z",
     "iopub.status.busy": "2023-10-09T16:24:33.073644Z",
     "iopub.status.idle": "2023-10-09T16:24:33.080798Z",
     "shell.execute_reply": "2023-10-09T16:24:33.079824Z"
    },
    "papermill": {
     "duration": 0.017363,
     "end_time": "2023-10-09T16:24:33.083009",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.065646",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  101,  3413,  8401,  4797,     1, 22724,  7058,   103, 24223,  9199,\n",
       "         11515,   488, 27828,  3547, 18158, 10189,     0,     0,     0,     0],\n",
       "        [  101,  4303, 15277,     1, 26430, 21759, 24603,   103, 12689, 15363,\n",
       "         14318,     5,  9876,     1,  4058,   451,     0,     0,     0,     0],\n",
       "        [  101, 12024, 27899,  8706,     1, 28294, 16359,   103,   835, 13595,\n",
       "         20214,     1,  8276, 11985, 14855,  8413,     0,     0,     0,     0],\n",
       "        [  101, 15900,   527, 15302, 26235, 16148, 10709,   103, 21784, 18223,\n",
       "         13851, 11017,     1, 25382,     1,     1,     0,     0,     0,     0]])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "defaee53",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.097864Z",
     "iopub.status.busy": "2023-10-09T16:24:33.097280Z",
     "iopub.status.idle": "2023-10-09T16:24:33.106570Z",
     "shell.execute_reply": "2023-10-09T16:24:33.105387Z"
    },
    "papermill": {
     "duration": 0.019335,
     "end_time": "2023-10-09T16:24:33.108740",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.089405",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-100, -100, -100, -100,    1, -100, -100, -100, -100, -100, -100, -100,\n",
       "         -100, -100, -100, -100, -100, -100, -100, -100],\n",
       "        [-100, -100, -100,    1, -100, -100, -100, -100, -100, -100, -100, -100,\n",
       "         -100,    1, -100, -100, -100, -100, -100, -100],\n",
       "        [-100, -100, -100, -100,    1, -100, -100, -100, -100, -100, -100,    1,\n",
       "         -100, -100, -100, -100, -100, -100, -100, -100],\n",
       "        [-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
       "            1, -100,    1,    1, -100, -100, -100, -100]])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = input_ids.clone()\n",
    "labels[masked_tokens==0]=-100\n",
    "labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "1310f5db",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.123820Z",
     "iopub.status.busy": "2023-10-09T16:24:33.123457Z",
     "iopub.status.idle": "2023-10-09T16:24:33.129850Z",
     "shell.execute_reply": "2023-10-09T16:24:33.129145Z"
    },
    "papermill": {
     "duration": 0.016274,
     "end_time": "2023-10-09T16:24:33.131841",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.115567",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.int64, torch.int64)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids.dtype, labels.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "f5371b7c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.148147Z",
     "iopub.status.busy": "2023-10-09T16:24:33.147520Z",
     "iopub.status.idle": "2023-10-09T16:24:33.308798Z",
     "shell.execute_reply": "2023-10-09T16:24:33.308022Z"
    },
    "papermill": {
     "duration": 0.172529,
     "end_time": "2023-10-09T16:24:33.311280",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.138751",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "head = nn.Linear(512,30_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "bcbf9aee",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.327868Z",
     "iopub.status.busy": "2023-10-09T16:24:33.326706Z",
     "iopub.status.idle": "2023-10-09T16:24:33.332189Z",
     "shell.execute_reply": "2023-10-09T16:24:33.331152Z"
    },
    "papermill": {
     "duration": 0.016502,
     "end_time": "2023-10-09T16:24:33.334631",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.318129",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "out = torch.rand(4,20,512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "8fc74146",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.350922Z",
     "iopub.status.busy": "2023-10-09T16:24:33.350273Z",
     "iopub.status.idle": "2023-10-09T16:24:33.416428Z",
     "shell.execute_reply": "2023-10-09T16:24:33.415393Z"
    },
    "papermill": {
     "duration": 0.076824,
     "end_time": "2023-10-09T16:24:33.418973",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.342149",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 20, 30000])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits = head(out)\n",
    "logits.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "75146ba1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.434021Z",
     "iopub.status.busy": "2023-10-09T16:24:33.433501Z",
     "iopub.status.idle": "2023-10-09T16:24:33.461179Z",
     "shell.execute_reply": "2023-10-09T16:24:33.459657Z"
    },
    "papermill": {
     "duration": 0.038572,
     "end_time": "2023-10-09T16:24:33.464256",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.425684",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "loss = F.cross_entropy(logits.view(-1,logits.size(-1)),labels.view(-1),ignore_index=-100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "7a8f47a9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.481964Z",
     "iopub.status.busy": "2023-10-09T16:24:33.481541Z",
     "iopub.status.idle": "2023-10-09T16:24:33.504779Z",
     "shell.execute_reply": "2023-10-09T16:24:33.504049Z"
    },
    "papermill": {
     "duration": 0.035185,
     "end_time": "2023-10-09T16:24:33.506846",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.471661",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(10.5416, grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "10338355",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.522549Z",
     "iopub.status.busy": "2023-10-09T16:24:33.522010Z",
     "iopub.status.idle": "2023-10-09T16:24:33.530606Z",
     "shell.execute_reply": "2023-10-09T16:24:33.529899Z"
    },
    "papermill": {
     "duration": 0.019239,
     "end_time": "2023-10-09T16:24:33.533032",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.513793",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class MultiheadAttention(nn.Module):\n",
    "    def __init__(self, dim, n_heads, dropout=0.):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.n_heads = n_heads\n",
    "        assert dim % n_heads == 0, 'dim should be div by n_heads'\n",
    "        self.head_dim = self.dim // self.n_heads\n",
    "        self.in_proj = nn.Linear(dim,dim*3)\n",
    "        self.attn_dropout = nn.Dropout(dropout)\n",
    "        self.scale = self.head_dim ** -0.5\n",
    "        self.out_proj = nn.Linear(dim,dim)\n",
    "        \n",
    "    def forward(self,x,mask=None):\n",
    "        b,t,c = x.shape\n",
    "        q,k,v = self.in_proj(x).chunk(3,dim=-1)\n",
    "        q = q.view(b,t,self.n_heads,self.head_dim).permute(0,2,1,3)\n",
    "        k = k.view(b,t,self.n_heads,self.head_dim).permute(0,2,1,3)\n",
    "        v = v.view(b,t,self.n_heads,self.head_dim).permute(0,2,1,3)\n",
    "        \n",
    "        qkT = torch.matmul(q,k.transpose(-1,-2)) * self.scale\n",
    "        qkT = self.attn_dropout(qkT)\n",
    "        \n",
    "        if mask is not None:\n",
    "            mask = mask.to(dtype=qkT.dtype,device=qkT.device)\n",
    "            qkT = qkT.masked_fill(mask==0,float('-inf'))\n",
    "            \n",
    "        qkT = F.softmax(qkT,dim=-1)\n",
    "            \n",
    "        attn = torch.matmul(qkT,v)\n",
    "        attn = attn.permute(0,2,1,3).contiguous().view(b,t,c)\n",
    "        out = self.out_proj(attn)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "362c7012",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.551362Z",
     "iopub.status.busy": "2023-10-09T16:24:33.550416Z",
     "iopub.status.idle": "2023-10-09T16:24:33.555754Z",
     "shell.execute_reply": "2023-10-09T16:24:33.555063Z"
    },
    "papermill": {
     "duration": 0.017541,
     "end_time": "2023-10-09T16:24:33.558042",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.540501",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self,dim,dropout=0.):\n",
    "        super().__init__()\n",
    "        self.feed_forward = nn.Sequential(\n",
    "            nn.Linear(dim,dim*4),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(dim*4,dim)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.feed_forward(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "98a8d6d4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.574791Z",
     "iopub.status.busy": "2023-10-09T16:24:33.574423Z",
     "iopub.status.idle": "2023-10-09T16:24:33.581299Z",
     "shell.execute_reply": "2023-10-09T16:24:33.579914Z"
    },
    "papermill": {
     "duration": 0.018228,
     "end_time": "2023-10-09T16:24:33.583459",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.565231",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    def __init__(self, dim, n_heads, attn_dropout=0., mlp_dropout=0.):\n",
    "        super().__init__()\n",
    "        self.attn = MultiheadAttention(dim,n_heads,attn_dropout)\n",
    "        self.ffd = FeedForward(dim,mlp_dropout)\n",
    "        self.ln_1 = nn.LayerNorm(dim)\n",
    "        self.ln_2 = nn.LayerNorm(dim)\n",
    "        \n",
    "    def forward(self,x,mask=None):\n",
    "        x = self.ln_1(x)\n",
    "        x = x + self.attn(x,mask)\n",
    "        x = self.ln_2(x)\n",
    "        x = x + self.ffd(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "6c10e51b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.599662Z",
     "iopub.status.busy": "2023-10-09T16:24:33.599000Z",
     "iopub.status.idle": "2023-10-09T16:24:33.604384Z",
     "shell.execute_reply": "2023-10-09T16:24:33.603545Z"
    },
    "papermill": {
     "duration": 0.016298,
     "end_time": "2023-10-09T16:24:33.606772",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.590474",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Embedding(nn.Module):\n",
    "    def __init__(self,vocab_size,max_len,dim):\n",
    "        super().__init__()\n",
    "        self.max_len = max_len\n",
    "        self.class_embedding = nn.Embedding(vocab_size,dim)\n",
    "        self.pos_embedding = nn.Embedding(max_len,dim)\n",
    "    def forward(self,x):\n",
    "        x = self.class_embedding(x)\n",
    "        pos = torch.arange(0,x.size(1),device=x.device)\n",
    "        x = x + self.pos_embedding(pos)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "a2316ed3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.624064Z",
     "iopub.status.busy": "2023-10-09T16:24:33.623691Z",
     "iopub.status.idle": "2023-10-09T16:24:33.633914Z",
     "shell.execute_reply": "2023-10-09T16:24:33.632360Z"
    },
    "papermill": {
     "duration": 0.022114,
     "end_time": "2023-10-09T16:24:33.636156",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.614042",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class MLMBERT(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.embedding = Embedding(config['vocab_size'],config['max_len'],config['dim'])\n",
    "        \n",
    "        self.depth = config['depth']\n",
    "        self.encoders = nn.ModuleList([\n",
    "            EncoderBlock(\n",
    "                dim=config['dim'],\n",
    "                n_heads=config['n_heads'],\n",
    "                attn_dropout=config['attn_dropout'],\n",
    "                mlp_dropout=config['mlp_dropout']\n",
    "            ) for _ in range(self.depth)\n",
    "        ])\n",
    "        \n",
    "        self.mlm_head = nn.Linear(config['dim'],config['vocab_size'])\n",
    "        \n",
    "        self.pad_token_id = config['pad_token_id']\n",
    "        self.mask_token_id = config['mask_token_id']\n",
    "        \n",
    "    def create_src_mask(self,src):\n",
    "        return (src != self.pad_token_id).unsqueeze(1).unsqueeze(2) # N, 1, 1, src_len\n",
    "    \n",
    "    def forward(self,input_ids,labels=None):\n",
    "        \n",
    "        src_mask = self.create_src_mask(input_ids)\n",
    "        enc_out = self.embedding(input_ids)\n",
    "        for layer in self.encoders:\n",
    "            enc_out = layer(enc_out,mask=src_mask)\n",
    "                \n",
    "        print(enc_out.shape)\n",
    "        logits = self.mlm_head(enc_out)\n",
    "        \n",
    "        if labels is not None:\n",
    "            loss = F.cross_entropy(logits.view(-1,logits.size(-1)),labels.view(-1))\n",
    "            return {'loss': loss, 'logits': logits}\n",
    "        else:\n",
    "            # assuming inference input_ids only have 1 [MASK] token\n",
    "            mask_idx = (input_ids==self.mask_token_id).flatten().nonzero().item()\n",
    "            mask_preds = F.softmax(logits[:,mask_idx,:],dim=-1).argmax(dim=-1)\n",
    "            return {'mask_predictions':mask_preds}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "c60faf92",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.654290Z",
     "iopub.status.busy": "2023-10-09T16:24:33.653828Z",
     "iopub.status.idle": "2023-10-09T16:24:33.660014Z",
     "shell.execute_reply": "2023-10-09T16:24:33.658553Z"
    },
    "papermill": {
     "duration": 0.017447,
     "end_time": "2023-10-09T16:24:33.662492",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.645045",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "config = {\n",
    "    'dim': 512,\n",
    "    'n_heads': 8,\n",
    "    'attn_dropout': 0.1,\n",
    "    'mlp_dropout': 0.1,\n",
    "    'depth': 6,\n",
    "    'vocab_size': 5_000,\n",
    "    'max_len': 128,\n",
    "    'pad_token_id': 0,\n",
    "    'mask_token_id': 1\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "d74b0775",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.679745Z",
     "iopub.status.busy": "2023-10-09T16:24:33.678698Z",
     "iopub.status.idle": "2023-10-09T16:24:33.686777Z",
     "shell.execute_reply": "2023-10-09T16:24:33.685519Z"
    },
    "papermill": {
     "duration": 0.019843,
     "end_time": "2023-10-09T16:24:33.689346",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.669503",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def sample():\n",
    "    input_ids = torch.randint(0,config['vocab_size'],(4,20))\n",
    "    input_ids[:,16:]=config['pad_token_id'] # PAD TOKEN\n",
    "    mlm_mask = torch.rand(input_ids.size()) < 0.15 * (input_ids!=config['pad_token_id'])\n",
    "    masked_tokens = input_ids * mlm_mask\n",
    "    labels = input_ids.clone()\n",
    "    labels[masked_tokens==config['pad_token_id']]=-100\n",
    "    input_ids[masked_tokens!=0]=config['mask_token_id'] # MASK TOKEN\n",
    "    return input_ids, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "36c5c65e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.706482Z",
     "iopub.status.busy": "2023-10-09T16:24:33.705567Z",
     "iopub.status.idle": "2023-10-09T16:24:33.713352Z",
     "shell.execute_reply": "2023-10-09T16:24:33.712308Z"
    },
    "papermill": {
     "duration": 0.018518,
     "end_time": "2023-10-09T16:24:33.715517",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.696999",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[2329,    1, 2274, 3112,    1,    1,    1, 4928,  321, 1492, 4794, 4964,\n",
      "            1, 2607, 1656, 4623,    0,    0,    0,    0],\n",
      "        [   1,    1, 3197, 2275, 2464, 2552, 4599,  269,  774,  721,  586,    1,\n",
      "         1804,    1,    1, 3224,    0,    0,    0,    0],\n",
      "        [2481,    1,  456, 1489, 1961,    1,    1, 3021,    1, 4425, 4870,    1,\n",
      "         3808, 2035, 4949, 4294,    0,    0,    0,    0],\n",
      "        [1893, 2362, 1262,  876, 1569,  718,    1, 3085, 1696,  775,    1, 4555,\n",
      "         2769, 2454, 2614,    1,    0,    0,    0,    0]])\n",
      "tensor([[-100, 3525, -100, -100, 4747, 4705, 3834, -100, -100, -100, -100, -100,\n",
      "         4951, -100, -100, -100, -100, -100, -100, -100],\n",
      "        [ 968,  424, -100, -100, -100, -100, -100, -100, -100, -100, -100, 1195,\n",
      "         -100, 1442, 4918, -100, -100, -100, -100, -100],\n",
      "        [-100, 3001, -100, -100, -100, 4988, 2002, -100, 4265, -100, -100,  467,\n",
      "         -100, -100, -100, -100, -100, -100, -100, -100],\n",
      "        [-100, -100, -100, -100, -100, -100, 1395, -100, -100, -100, 2625, -100,\n",
      "         -100, -100, -100, 1313, -100, -100, -100, -100]])\n"
     ]
    }
   ],
   "source": [
    "i,l = sample()\n",
    "print(i)\n",
    "print(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "0f93ae8d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.730376Z",
     "iopub.status.busy": "2023-10-09T16:24:33.729787Z",
     "iopub.status.idle": "2023-10-09T16:24:33.735923Z",
     "shell.execute_reply": "2023-10-09T16:24:33.735162Z"
    },
    "papermill": {
     "duration": 0.015804,
     "end_time": "2023-10-09T16:24:33.737708",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.721904",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([4, 20]), torch.Size([4, 20]))"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i.shape, l.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "32f0e63c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:33.752329Z",
     "iopub.status.busy": "2023-10-09T16:24:33.751967Z",
     "iopub.status.idle": "2023-10-09T16:24:34.020523Z",
     "shell.execute_reply": "2023-10-09T16:24:34.019381Z"
    },
    "papermill": {
     "duration": 0.278787,
     "end_time": "2023-10-09T16:24:34.023053",
     "exception": false,
     "start_time": "2023-10-09T16:24:33.744266",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model = MLMBERT(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "25786540",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:34.038534Z",
     "iopub.status.busy": "2023-10-09T16:24:34.037925Z",
     "iopub.status.idle": "2023-10-09T16:24:34.193834Z",
     "shell.execute_reply": "2023-10-09T16:24:34.192166Z"
    },
    "papermill": {
     "duration": 0.167296,
     "end_time": "2023-10-09T16:24:34.196915",
     "exception": false,
     "start_time": "2023-10-09T16:24:34.029619",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 20, 512])\n"
     ]
    }
   ],
   "source": [
    "out = model(i,l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "286f9bcb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-09T16:24:34.215250Z",
     "iopub.status.busy": "2023-10-09T16:24:34.214795Z",
     "iopub.status.idle": "2023-10-09T16:24:34.222014Z",
     "shell.execute_reply": "2023-10-09T16:24:34.220641Z"
    },
    "papermill": {
     "duration": 0.01914,
     "end_time": "2023-10-09T16:24:34.224181",
     "exception": false,
     "start_time": "2023-10-09T16:24:34.205041",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 20, 5000])"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out['logits'].shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 9.359797,
   "end_time": "2023-10-09T16:24:35.308229",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2023-10-09T16:24:25.948432",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}