Model Definations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model \n",
    "This notebook is used for designing the custom Albert model. The model is also being tested as sanity check. The custom model is implemented as a Pytorch Lightning module.\n",
    "\n",
    "**Current Tasks Remaining**:-\n",
    "1. Implement the training_step(), configure_optimizer().\n",
    "2. Modify the dataloaders. \n",
    "3. sanity Check run.\n",
    "4. Increase MLM data size.\n",
    "\n",
    "The forward function needs to have all the parameters mentioned at this [link](https://github.com/huggingface/transformers/blob/master/src/transformers/models/albert/modeling_albert.py#L640). The output from this model will then be passed into the MLM head of the Albert and the loss used will be **CrossEntropyLoss** with **Adam** Optimizer.\n",
    "\n",
    "### Note:\n",
    "Will also need to check the batch size permissible along with the gpus option on PL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AlbertConfig, AlbertModel\n",
    "from transformers.modeling_albert import AlbertMLMHead\n",
    "import torch\n",
    "from torch import nn\n",
    "import pytorch_lightning as pl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining custom Albert Config\n",
    "\n",
    "Refer to [this](https://github.com/huggingface/transformers/blob/48cc224703a8dd8d03d2721c8651fea8704d994b/src/transformers/models/albert/configuration_albert.py#L33) link to understand the meaning of the parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters for our model\n",
    "vocab_size = 10\n",
    "embedding_size = 16\n",
    "hidden_size = 768\n",
    "num_attention_heads = 12\n",
    "intermediate_size = 3072\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining Custom Albert Model config\n",
    "custom_config = AlbertConfig(\n",
    "    vocab_size=vocab_size, # A, C, T, G, U, UNK, MASK, PAD, CLS, SEP\n",
    "    embedding_size=embedding_size, #this will be scaled to 32 and 64 for ablation experiments\n",
    "    hidden_size=hidden_size,\n",
    "    num_attention_heads=num_attention_heads,\n",
    "    intermediate_size=intermediate_size,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_model = AlbertModel(custom_config) # custom model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#downloading from huggingface the pretrained model\n",
    "pretrained_model = AlbertModel.from_pretrained('albert-base-v2', return_dict=True)\n",
    "# pretrained_model.save_pretrained('./albert_base_v2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# masked language modelling head\n",
    "# this MLM head will be put on top our custom config albert\n",
    "# and trained on miRNA and mRNA sequences separately.\n",
    "mlm_head = AlbertMLMHead(custom_config) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AlbertModel(\n",
       "  (embeddings): AlbertEmbeddings(\n",
       "    (word_embeddings): Embedding(30000, 128, padding_idx=0)\n",
       "    (position_embeddings): Embedding(512, 128)\n",
       "    (token_type_embeddings): Embedding(2, 128)\n",
       "    (LayerNorm): LayerNorm((128,), eps=1e-12, elementwise_affine=True)\n",
       "    (dropout): Dropout(p=0, inplace=False)\n",
       "  )\n",
       "  (encoder): AlbertTransformer(\n",
       "    (embedding_hidden_mapping_in): Linear(in_features=128, out_features=768, bias=True)\n",
       "    (albert_layer_groups): ModuleList(\n",
       "      (0): AlbertLayerGroup(\n",
       "        (albert_layers): ModuleList(\n",
       "          (0): AlbertLayer(\n",
       "            (full_layer_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (attention): AlbertAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (attention_dropout): Dropout(p=0, inplace=False)\n",
       "              (output_dropout): Dropout(p=0, inplace=False)\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            )\n",
       "            (ffn): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (ffn_output): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (pooler): Linear(in_features=768, out_features=768, bias=True)\n",
       "  (pooler_activation): Tanh()\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pretrained_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AlbertModel(\n",
       "  (embeddings): AlbertEmbeddings(\n",
       "    (word_embeddings): Embedding(10, 16, padding_idx=0)\n",
       "    (position_embeddings): Embedding(512, 16)\n",
       "    (token_type_embeddings): Embedding(2, 16)\n",
       "    (LayerNorm): LayerNorm((16,), eps=1e-12, elementwise_affine=True)\n",
       "    (dropout): Dropout(p=0, inplace=False)\n",
       "  )\n",
       "  (encoder): AlbertTransformer(\n",
       "    (embedding_hidden_mapping_in): Linear(in_features=16, out_features=768, bias=True)\n",
       "    (albert_layer_groups): ModuleList(\n",
       "      (0): AlbertLayerGroup(\n",
       "        (albert_layers): ModuleList(\n",
       "          (0): AlbertLayer(\n",
       "            (full_layer_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (attention): AlbertAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (attention_dropout): Dropout(p=0, inplace=False)\n",
       "              (output_dropout): Dropout(p=0, inplace=False)\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            )\n",
       "            (ffn): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (ffn_output): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (pooler): Linear(in_features=768, out_features=768, bias=True)\n",
       "  (pooler_activation): Tanh()\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since the pretrained model has a vocabulary of 30k\n",
    "# and embedding size of 128, we need to downscale it\n",
    "# to our requirements\n",
    "# Thus, here the embeddings are not pretrained but\n",
    "# only the main model is. Objective of this is to \n",
    "# leverage the latent space of pretrained model.\n",
    "\n",
    "# pretrained_model.resize_token_embeddings(10)\n",
    "# pretrained_model.set_input_embeddings(nn.Embedding(10,16, padding_idx=0))\n",
    "# pretrained_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AlbertMLMHead(\n",
       "  (LayerNorm): LayerNorm((16,), eps=1e-05, elementwise_affine=True)\n",
       "  (dense): Linear(in_features=768, out_features=16, bias=True)\n",
       "  (decoder): Linear(in_features=16, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mlm_head"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating Text Dataset and Dataloader\n",
    "In this section, we create the **Dataset** and **Dataloader** to be used in our training and associated tasks. Since our vocabulary is limited, we are not using any packages like ***SpaCy***. We create a vocabulary set and our own tokenizer for our models.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vocabulary:\n",
    "    def __init__(self):\n",
    "\n",
    "        self.itos = {0: \"[PAD]\", 1: \"[START]\", 2: \"[END]\", 3: \"[UNK]\", 4: \"[MASK]\"}\n",
    "        self.stoi = {\"[PAD]\": 0, \"[START]\": 1, \"[END]\": 2, \"[UNK]\": 3, \"[MASK]\": 4}\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.itos)\n",
    "\n",
    "    #returns the numeric token value of a given string token\n",
    "    def get_idx(self, token):\n",
    "        return self.stoi[token]\n",
    "    \n",
    "    #returns the alphanumeric token value given the token idx\n",
    "    def get_token(self, token):\n",
    "        return self.itos[token]\n",
    "\n",
    "    @staticmethod\n",
    "    def tokenizer_seq(fasta_seq):\n",
    "#         print(fasta_seq)\n",
    "        return [str(x) for x in list(fasta_seq)]\n",
    "\n",
    "    def build_vocabulary(self):\n",
    "        frequencies = {}\n",
    "        idx = len(self.itos)\n",
    "        for idx1, base in enumerate(list('acgut')):\n",
    "            self.stoi[base] = idx+idx1\n",
    "            self.itos[idx+idx1] = base\n",
    "\n",
    "    def numericalize(self, fasta_seq):\n",
    "        tokenized_seq = self.tokenizer_seq(fasta_seq.lower())\n",
    "        return [\n",
    "            self.stoi[token] if token in self.stoi else self.stoi[\"[UNK]\"]\n",
    "            for token in tokenized_seq\n",
    "        ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SequenceDataset(Dataset):\n",
    "    def __init__(self, filename, sep='\\t'):\n",
    "        self.df = pd.read_csv(filename,sep=sep, index_col=0)\n",
    "#         print(self.df.head())\n",
    "        # Dataset Column Positions - mRNA  Binding sites with flanking regions\n",
    "        self.mirna_names = self.df.iloc[:, 0].values\n",
    "        self.mirna_seqs = self.df.iloc[:, 1].values\n",
    "        #concatenating row-wise to create a combined vocabulary\n",
    "#         all_seq = self.mirna[:] + self.mrna\n",
    "        \n",
    "        # Initialize vocabulary and build vocab\n",
    "        self.vocab = Vocabulary()\n",
    "        self.vocab.build_vocabulary()\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.df.shape[0]\n",
    "\n",
    "    def numericalize_seq(self,seq):\n",
    "        numericalized_seq = [self.vocab.stoi[\"[START]\"]]\n",
    "        numericalized_seq += self.vocab.numericalize(seq)\n",
    "        numericalized_seq.append(self.vocab.stoi[\"[END]\"])\n",
    "        return numericalized_seq\n",
    "\n",
    "    def get_vocabulary(self):\n",
    "        return self.vocab.stoi\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "#         print(index, self.mirna_names[index])\n",
    "        return torch.tensor(self.numericalize_seq(self.mirna_seqs[index]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CollateSequences:\n",
    "    def __init__(self, pad_idx):\n",
    "        self.pad_idx = pad_idx\n",
    "\n",
    "    def __call__(self, batch):\n",
    "        \n",
    "        seq = [item for item in batch]\n",
    "#         print(seq)\n",
    "        seq = pad_sequence(seq, batch_first=True, padding_value=self.pad_idx)\n",
    "\n",
    "        return seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/abhik/anaconda3/envs/pytorch_env/lib/python3.8/site-packages/numpy/lib/arraysetops.py:569: FutureWarning: elementwise comparison failed; returning scalar instead, but in the future will perform elementwise comparison\n",
      "  mask |= (ar1 == a)\n"
     ]
    }
   ],
   "source": [
    "batch_size=5\n",
    "num_workers=1\n",
    "shuffle=True\n",
    "pin_memory=True\n",
    "\n",
    "dataset = SequenceDataset(filename='./Processed MBStar/partitioned/mlm/mlm_mrna_data_chunk_0.txt')\n",
    "\n",
    "pad_idx = dataset.vocab.stoi[\"[PAD]\"]\n",
    "\n",
    "loader = DataLoader(\n",
    "        dataset=dataset,\n",
    "        batch_size=batch_size,\n",
    "        num_workers=num_workers,\n",
    "        shuffle=shuffle,\n",
    "        pin_memory=pin_memory,\n",
    "        collate_fn=CollateSequences(pad_idx=pad_idx)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.0098, -1.1421, -0.0952,  ..., -0.6524,  0.0876,  0.9149],\n",
      "         [ 0.0165, -1.1618, -0.2238,  ..., -0.7277, -0.1072,  0.8638],\n",
      "         [-0.0869, -1.1878, -0.1202,  ..., -0.7930, -0.0861,  1.0118],\n",
      "         ...,\n",
      "         [-0.1806, -1.2484, -0.0447,  ..., -0.6611,  0.1465,  1.0092],\n",
      "         [-0.0567, -1.1168, -0.3301,  ..., -0.6464, -0.0097,  0.8886],\n",
      "         [-0.1023, -0.9794, -0.2897,  ..., -0.6191,  0.0302,  0.9009]],\n",
      "\n",
      "        [[-0.0836, -1.1591, -0.0430,  ..., -0.6884,  0.0159,  1.0127],\n",
      "         [-0.2882, -0.9896, -0.2202,  ..., -0.8375, -0.3605,  1.0494],\n",
      "         [-0.1954, -1.2341, -0.1343,  ..., -0.9152, -0.2301,  1.1028],\n",
      "         ...,\n",
      "         [-0.2595, -1.3284,  0.3073,  ..., -0.7161, -0.0194,  1.1922],\n",
      "         [-0.3094, -1.2860,  0.3828,  ..., -0.8379, -0.3138,  1.1856],\n",
      "         [-0.1684, -0.9976, -0.4257,  ..., -0.6289, -0.0691,  0.9643]],\n",
      "\n",
      "        [[-0.1534, -1.0976,  0.0412,  ..., -0.7163, -0.1666,  1.0131],\n",
      "         [-0.1998, -1.0464, -0.0654,  ..., -0.8510, -0.3852,  1.0770],\n",
      "         [-0.3045, -1.1808,  0.2970,  ..., -0.7778, -0.2219,  1.2119],\n",
      "         ...,\n",
      "         [-0.2746, -1.2561,  0.0881,  ..., -0.7329,  0.0192,  1.0515],\n",
      "         [-0.1395, -1.1498, -0.3015,  ..., -0.7283, -0.2102,  0.9626],\n",
      "         [-0.1699, -1.0069, -0.2536,  ..., -0.6848, -0.1274,  0.9812]],\n",
      "\n",
      "        [[ 0.0606, -1.1711, -0.0748,  ..., -0.6349,  0.1103,  0.8334],\n",
      "         [ 0.1164, -1.1828, -0.1663,  ..., -0.6781, -0.0338,  0.7834],\n",
      "         [-0.1101, -1.1098, -0.2665,  ..., -0.6746,  0.1712,  0.9728],\n",
      "         ...,\n",
      "         [-0.0873, -1.3149,  0.2760,  ..., -0.6342,  0.1178,  1.0576],\n",
      "         [-0.0036, -1.1109, -0.3145,  ..., -0.6212,  0.0175,  0.8173],\n",
      "         [-0.0527, -0.9875, -0.2744,  ..., -0.6014,  0.0530,  0.8398]],\n",
      "\n",
      "        [[-0.0105, -1.1495, -0.0124,  ..., -0.6518,  0.0328,  0.9581],\n",
      "         [ 0.0015, -0.9971, -0.1838,  ..., -0.6287,  0.0615,  0.8178],\n",
      "         [-0.2165, -1.2217,  0.2458,  ..., -0.7504, -0.0808,  1.1593],\n",
      "         ...,\n",
      "         [-0.0436, -1.2664,  0.1180,  ..., -0.6523,  0.2686,  0.9472],\n",
      "         [-0.0588, -1.1654, -0.2878,  ..., -0.6678, -0.0500,  0.9162],\n",
      "         [-0.1041, -1.0163, -0.2427,  ..., -0.6362, -0.0027,  0.9340]]],\n",
      "       grad_fn=<AddBackward0>)\n"
     ]
    }
   ],
   "source": [
    "#testing dataloader\n",
    "import numpy as np\n",
    "for _, a in enumerate(loader):\n",
    "    print(a)\n",
    "    print(a.shape)\n",
    "    print(a.masked_fill(a != 0, value=1))\n",
    "    \n",
    "    print(logits)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Masked Language Modelling using Albert\n",
    "This section demos how to use the original Albert model with the original Tokenizer for masked language modelling task. The mask token used in this case is \\[MASK\\]. The model is tasked with predicting what this word might be. \n",
    "The model outputs logits and losses. The logits are at position 0 in the output. The key \\'last_hidden_state\\' can be used to extract them as well. This logit is sent through a **Masked Language Modelling Head** which outputs a probability for each of the words from the vocabulary to be placed instead of the \\[MASK\\]. The maximum from this output is the token for the word. This is found by using **torch.argmax** on the array elememt in the output at the position of the  \\[MASK\\] in the second dimension (output dimension is ***\\[batch_size, number_of_token_in_input, number_of_words_in_vocabulary\\]*** )\n",
    "\n",
    "#### Note:\n",
    "Since an untrained MLM HEAD is used, the output may be arbitrary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#loading the pretrained original tokenizer used for albert-base-v2\n",
    "from transformers import AlbertTokenizer\n",
    "from transformers import AlbertConfig, AlbertModel\n",
    "from transformers.modeling_albert import AlbertMLMHead\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AlbertMLMHead(\n",
       "  (LayerNorm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "  (dense): Linear(in_features=768, out_features=128, bias=True)\n",
       "  (decoder): Linear(in_features=128, out_features=30000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pretrained_model = AlbertModel.from_pretrained('albert-base-v2', return_dict=True) #pretrained model\n",
    "tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2') #tokenizer\n",
    "\n",
    "#This is the Masked Language Modelling HEAD for original Albert\n",
    "albert_base_configuration = AlbertConfig(\n",
    "      hidden_size=768,\n",
    "      num_attention_heads=12,\n",
    "      intermediate_size=3072,\n",
    "  )\n",
    "default_mlm_head = AlbertMLMHead(albert_base_configuration) #default mlm head\n",
    "default_mlm_head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encoded Sequence- [2, 48, 25, 21, 4, 669, 3]\n",
      "ID for [MASK] token - [MASK]\n",
      "[MASK] is at- 5\n"
     ]
    }
   ],
   "source": [
    "#change the sentance as long as it is within 100 words\n",
    "#use only 1 [MASK]\n",
    "sentence = 'This is a [MASK] project' \n",
    "encoded_seq = tokenizer.encode(sentence)\n",
    "print(\"Encoded Sequence-\", encoded_seq)\n",
    "print(\"ID for [MASK] token -\",tokenizer.convert_ids_to_tokens(4))\n",
    "mask_index = encoded_seq.index(4)\n",
    "print(\"[MASK] is at-\", mask_index+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Keys-  odict_keys(['last_hidden_state', 'pooler_output'])\n",
      "Sequence Output- tensor([[[ 0.8959,  0.7075,  0.7822,  ..., -0.5718,  0.8908,  0.2283],\n",
      "         [ 1.4882, -1.0936,  0.3678,  ..., -0.2871,  2.6546, -1.4538],\n",
      "         [ 1.3232, -0.9379, -0.5643,  ...,  0.7148,  0.4633, -0.5228],\n",
      "         ...,\n",
      "         [-0.4393, -0.3515, -1.2364,  ..., -0.2335,  0.6026, -1.3187],\n",
      "         [-0.3884, -0.5458,  0.5886,  ...,  0.2667,  0.8635, -0.4085],\n",
      "         [ 0.0666,  0.1381, -0.0645,  ..., -0.0783,  0.1365,  0.1967]]],\n",
      "       grad_fn=<NativeLayerNormBackward>)\n",
      "Sequence Output Size- torch.Size([1, 7, 768])\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Consult https://github.com/huggingface/transformers/blob/master/src/transformers/models/albert/modeling_albert.py#L698\n",
    "for understanding output format\n",
    "\"\"\"\n",
    "\n",
    "test_output = pretrained_model(torch.tensor(encoded_seq).unsqueeze(0))\n",
    "print(\"Keys- \",test_output.keys())\n",
    "print(\"Sequence Output-\",test_output.last_hidden_state)\n",
    "print(\"Sequence Output Size-\",test_output.last_hidden_state.size())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction Scores Dimensions- torch.Size([1, 7, 30000])\n",
      "Predicted Token for [MASK]- tensor(27047)\n"
     ]
    }
   ],
   "source": [
    "prediction_scores = default_mlm_head(test_output.last_hidden_state)\n",
    "print(\"Prediction Scores Dimensions-\",prediction_scores.size())\n",
    "print(\"Predicted Token for [MASK]-\",torch.argmax(prediction_scores[0, mask_index]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted -  ▁sykes\n"
     ]
    }
   ],
   "source": [
    "predicted_word = tokenizer.convert_ids_to_tokens(torch.argmax(prediction_scores[0, mask_index]).item())\n",
    "print(\"Predicted - \", predicted_word)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Designing Model\n",
    "In this section, the custom AlBERT model is being defined which will be trained on MLM on sequences. The model is a *PyTorch Lightning* model. We are going to use CrossEntropyLoss and Adam Optimizer with Weight Decay.\n",
    "Consult [this](https://huggingface.co/transformers/_modules/transformers/optimization.html#AdamW) and [this](https://huggingface.co/transformers/main_classes/optimizer_schedules.html) for AdamW and Optimization Schedules available through HuggingFace Transformers respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AlbertConfig, AlbertModel\n",
    "from transformers.modeling_albert import AlbertMLMHead\n",
    "from transformers import AdamW\n",
    "from torch.nn import CrossEntropyLoss\n",
    "import torch\n",
    "from torch import nn\n",
    "import numpy as np\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "class CustomAlbert(pl.LightningModule):\n",
    "    def __init__(\n",
    "        self,\n",
    "        vocab_size=10,\n",
    "        embedding_size=16,\n",
    "        hidden_size=768,\n",
    "        num_attention_heads=12,\n",
    "        intermediate_size=3072):\n",
    "        \n",
    "        super(CustomAlbert, self).__init__()\n",
    "        \n",
    "        self.save_hyperparameters()\n",
    "        \n",
    "        self.vocab_size = vocab_size\n",
    "        self.embedding_size = embedding_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_attention_heads = num_attention_heads\n",
    "        self.intermediate_size = intermediate_size\n",
    "        \n",
    "        custom_config = AlbertConfig(\n",
    "            vocab_size=vocab_size, # A, C, T, G, U, UNK, MASK, PAD, CLS, SEP\n",
    "            embedding_size=embedding_size, #this will be scaled to 32 and 64 for ablation experiments\n",
    "            hidden_size=hidden_size,\n",
    "            num_attention_heads=num_attention_heads,\n",
    "            intermediate_size=intermediate_size,\n",
    "        )\n",
    "        \n",
    "        self.custom_model = AlbertModel(custom_config) # custom model\n",
    "        self.mlm_head = AlbertMLMHead(custom_config) # mlm head\n",
    "        self.loss = CrossEntropyLoss()\n",
    "        \n",
    "\n",
    "    def forward(\n",
    "        self, \n",
    "        input_ids=None,\n",
    "        attention_mask=None,\n",
    "        token_type_ids=None,\n",
    "        position_ids=None,\n",
    "        head_mask=None,\n",
    "        inputs_embeds=None,\n",
    "        labels=None,\n",
    "        sentence_order_label=None,\n",
    "        output_attentions=None,\n",
    "        output_hidden_states=None,\n",
    "        return_dict=None):\n",
    "        \n",
    "        outputs = self.custom_model(\n",
    "            input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            token_type_ids=token_type_ids,\n",
    "            position_ids=position_ids,\n",
    "            head_mask=head_mask,\n",
    "            inputs_embeds=inputs_embeds,\n",
    "            output_attentions=output_attentions,\n",
    "            output_hidden_states=output_hidden_states,\n",
    "            return_dict=return_dict,\n",
    "        )\n",
    "\n",
    "        sequence_output, pooled_output = outputs[:2]\n",
    "        \n",
    "        prediction_scores = self.mlm_head(sequence_output)\n",
    "        \n",
    "        return prediction_scores\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        self.optimizer = AdamW(self.parameters(), lr=1e-3) #default learning rate is 1e-3\n",
    "        \n",
    "        return self.optimizer\n",
    "    \n",
    "    def training_step(self, batch):\n",
    "        seq = batch\n",
    "        attention_mask = seq.masked_fill( seq != 0, value=1)\n",
    "        logits = self(input_ids=seq, attention_mask=attention_mask)\n",
    "        \n",
    "        # For loss calculation, only sequence-based loss and not the loss from pads \n",
    "        # is taken into account. This is found by using the indices where\n",
    "        # attention_mask element is 1.\n",
    "        active_loss = attention_mask.view(-1) == 1 \n",
    "        active_logits = logits.view(-1, self.num_labels)[active_loss] \n",
    "        active_labels = seq.view(-1)[active_loss]\n",
    "        \n",
    "        loss = self.loss(active_logits, active_labels)\n",
    "        \n",
    "        return loss\n",
    "        \n",
    "    def train_dataloader(self):\n",
    "        dataset = SequenceDataset(filename='./Processed MBStar/partitioned/mlm/mlm_mrna_data_chunk_0.txt')\n",
    "\n",
    "        pad_idx = dataset.vocab.stoi[\"[PAD]\"]\n",
    "\n",
    "        loader = DataLoader(\n",
    "                dataset=dataset,\n",
    "                batch_size=batch_size,\n",
    "                num_workers=num_workers,\n",
    "                shuffle=shuffle,\n",
    "                pin_memory=pin_memory,\n",
    "                collate_fn=CollateSequences(pad_idx=pad_idx)\n",
    "            )\n",
    "        \n",
    "        return loader\n",
    "        "
   ]
  },
  {
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}