Attention_is_all_you_need.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 64,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iZfld1v1iEmf",
        "outputId": "ec0eb954-0035-414c-e3ff-d8717710f3db"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "cuda\n"
          ]
        }
      ],
      "source": [
        "from torch.utils.data import Dataset\n",
        "import torch.nn as nn\n",
        "import numpy as np\n",
        "np.random.seed(42)\n",
        "import torch\n",
        "import json\n",
        "from pickle import load\n",
        "import torchtext\n",
        "import torch\n",
        "import numpy as np\n",
        "from torch.nn import CrossEntropyLoss\n",
        "device='cuda' if torch.cuda.is_available() else 'cpu'\n",
        "print(device)\n",
        "from torch.optim import Adam"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## <font color='red'> Custom Tokenizer to tokenize the text data </font>"
      ],
      "metadata": {
        "id": "twrgzGxWi4Mr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class CustomTokenizer:\n",
        "    def __init__(self, dataset=None, vocab=None):\n",
        "        # Initialize the CustomTokenizer with a dataset and a custom vocabulary (if provided)\n",
        "\n",
        "        # Create a custom vocabulary from the dataset if no vocabulary is provided\n",
        "        if vocab is None:\n",
        "            self.vocab = {}\n",
        "            self.pad_index = 0\n",
        "            self.vocab['<pad>'] = self.pad_index\n",
        "            self.sos_index = len(self.vocab)\n",
        "            self.vocab['<sos>'] = self.sos_index\n",
        "            self.eos_index = len(self.vocab)\n",
        "            self.vocab['<eos>'] = self.eos_index\n",
        "\n",
        "            # Iterate through the dataset to build the vocabulary\n",
        "            for text in dataset:\n",
        "                # Exclude the first and last 5 characters from each text\n",
        "                text = text[5:-5]\n",
        "\n",
        "                # Tokenize the text and add tokens to the vocabulary\n",
        "                for token in text.split():\n",
        "                    if token not in self.vocab:\n",
        "                        self.vocab[token] = len(self.vocab)\n",
        "        else:\n",
        "            # Use the provided vocabulary\n",
        "            self.vocab = vocab\n",
        "\n",
        "    def get_vocab(self):\n",
        "        # Get the custom vocabulary\n",
        "        return self.vocab\n",
        "\n",
        "    def tokenize(self, text, max_length):\n",
        "        # Tokenize a text using the custom vocabulary\n",
        "\n",
        "        # Exclude the first and last 5 characters from the text\n",
        "        text = text[5:-5]\n",
        "\n",
        "        # Initialize the encoded vector with the '<sos>' token\n",
        "        encoded_vector = [self.vocab['<sos>']]\n",
        "\n",
        "        # Remove non-alphanumeric characters from the text\n",
        "        text = ''.join([i for i in text if i.isalnum() or i == ' '])\n",
        "\n",
        "        # Tokenize the text and add tokens to the encoded vector\n",
        "        for i, token in enumerate(text.split()):\n",
        "            if i >= max_length - 2:\n",
        "                break\n",
        "            if token in self.vocab:\n",
        "                encoded_vector.append(self.vocab[token])\n",
        "\n",
        "        # Add the '<eos>' token to the end of the encoded vector\n",
        "        encoded_vector.append(self.vocab['<eos>'])\n",
        "\n",
        "        # Add padding to the encoded vector if its length is less than max_length\n",
        "        if len(encoded_vector) < max_length:\n",
        "            # Add padding till max_length\n",
        "            for i in range(max_length - len(encoded_vector)):\n",
        "                encoded_vector.append(self.vocab['<pad>'])\n",
        "\n",
        "        # Return the encoded vector as a PyTorch tensor\n",
        "        return torch.tensor(encoded_vector)\n",
        "\n",
        "    def tokenize_batch(self, text_list, max_len):\n",
        "        # Tokenize a list of texts using the custom vocabulary\n",
        "        encoded_batch = torch.stack([self.tokenize(text, max_len) for text in text_list])\n",
        "        return encoded_batch\n",
        "\n",
        "    def get_vocab_size(self):\n",
        "        # Get the size of the custom vocabulary\n",
        "        return len(self.vocab)\n"
      ],
      "metadata": {
        "id": "e6ImTI9AiZpQ"
      },
      "execution_count": 65,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## <font color='red'> Dataloader class to load the data, split into train and test </font>"
      ],
      "metadata": {
        "id": "Mej9HFWnjKxE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TranslationDataloader(Dataset):\n",
        "    def __init__(self,filepath,max_sentences):\n",
        "        super().__init__()\n",
        "        english_data,german_data=self.load_data_from_file(filepath)\n",
        "        self.english_data=english_data[:max_sentences]\n",
        "        self.german_data=german_data[:max_sentences]\n",
        "        self.shuffle_indices=np.arange(len(self.english_data))\n",
        "        np.random.shuffle(self.shuffle_indices)\n",
        "        self.train_split=0.85\n",
        "        self.test_split=0.15\n",
        "        self.eng_train_data,self.ger_train_data,self.eng_test_data,self.ger_test_data=self.get_train_test_split()\n",
        "\n",
        "        print(\"Train Data Size: \",len(self.eng_train_data))\n",
        "        print(\"Test Data Size: \",len(self.eng_test_data))\n",
        "        english_vocab=CustomTokenizer(dataset=self.eng_train_data)\n",
        "        german_vocab=CustomTokenizer(dataset=self.ger_train_data)\n",
        "        self.eng_tokenizer=english_vocab.get_vocab()\n",
        "        self.ger_tokenizer=german_vocab.get_vocab()\n",
        "\n",
        "        #get maximum length of german sentences\n",
        "        # self.max_german_length=max([len(sentence.split()) for sentence in self.german_data])\n",
        "        # self.max_english_length=max([len(sentence.split()) for sentence in self.english_data])\n",
        "        self.max_german_length=33\n",
        "        self.max_english_length=33\n",
        "        print(\"Max German Sentence Length: \",self.max_german_length)\n",
        "        print(\"Max English Sentence Length: \",self.max_english_length)\n",
        "        #save eng and ger tokenizer\n",
        "        with open('eng_tokenizer.json', 'w') as fp:\n",
        "            json.dump(self.eng_tokenizer, fp)\n",
        "        with open('ger_tokenizer.json', 'w') as fp:\n",
        "            json.dump(self.ger_tokenizer, fp)\n",
        "        print(\"English Vocab Size: \",len(self.eng_tokenizer))\n",
        "        print(\"German Vocab Size: \",len(self.ger_tokenizer))\n",
        "\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.english_data)\n",
        "    def __getitem__(self,idx):\n",
        "        return self.english_data[idx],self.german_data[idx]\n",
        "\n",
        "    def load_data_from_file(self,filepath):\n",
        "        dataset=load(open(filepath, 'rb'))\n",
        "        english_sentences=['<sos>'+sentence[0]+'<eos>' for sentence in dataset]\n",
        "        german_sentences=['<sos>'+sentence[1]+'<eos>' for sentence in dataset]\n",
        "        return english_sentences,german_sentences\n",
        "\n",
        "\n",
        "    def get_train_test_split(self):\n",
        "        train_end_index=int(self.train_split*len(self.english_data))\n",
        "        print(train_end_index)\n",
        "        print()\n",
        "        eng_train_data=self.english_data[:train_end_index]\n",
        "        ger_train_data=self.german_data[:train_end_index]\n",
        "        eng_test_data=self.english_data[train_end_index:]\n",
        "        ger_test_data=self.german_data[train_end_index:]\n",
        "        return eng_train_data,ger_train_data,eng_test_data,ger_test_data\n",
        "\n",
        "    def create_encoded_data(self,data,vocab,max_len):\n",
        "        encoded_data=torch.zeros((len(data),max_len),dtype=torch.long)\n",
        "        for index,text in enumerate(data):\n",
        "\n",
        "            encoded_vector=self.create_encoded_vector(text,vocab,max_len)\n",
        "            encoded_data[index]=torch.tensor(encoded_vector,dtype=torch.long)\n",
        "\n",
        "        return encoded_data\n",
        "\n",
        "    def create_encoded_vector(self,text,vocab,max_len):\n",
        "            text=text[5:-5]\n",
        "            encoded_vector=[]\n",
        "            encoded_vector.append(vocab['<sos>'])\n",
        "            for i,token in enumerate(text.split()):\n",
        "                if(i>=max_len-2):\n",
        "                    break\n",
        "                if token in vocab:\n",
        "                    encoded_vector.append(vocab[token])\n",
        "\n",
        "            encoded_vector.append(vocab['<eos>'])\n",
        "            for k in range(max_len-len(encoded_vector)):\n",
        "                encoded_vector.append(vocab['<pad>'])\n",
        "            return encoded_vector\n",
        "\n",
        "\n",
        "    def get_encoded_data(self):\n",
        "        train_encoded_eng=self.create_encoded_data(self.eng_train_data,self.eng_tokenizer,self.max_english_length)\n",
        "        train_encoded_ger=self.create_encoded_data(self.ger_train_data,self.ger_tokenizer,self.max_german_length)\n",
        "        test_encoded_eng=self.create_encoded_data(self.eng_test_data,self.eng_tokenizer,self.max_english_length)\n",
        "        test_encoded_ger=self.create_encoded_data(self.ger_test_data,self.ger_tokenizer,self.max_german_length)\n",
        "        return train_encoded_eng,train_encoded_ger,test_encoded_eng,test_encoded_ger\n",
        ""
      ],
      "metadata": {
        "id": "U4rXu4y-iOo8"
      },
      "execution_count": 66,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## <font color='red'> Word and positional embeddings </font>"
      ],
      "metadata": {
        "id": "3W8NTN6qj9Lc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TokenPositionEmbeddings(nn.Module):\n",
        "    def __init__(self, vocab_size, max_len, embedding_dim):\n",
        "        # Initialize TokenPositionEmbeddings module with vocabulary size, maximum length, and embedding dimension\n",
        "\n",
        "        super().__init__()\n",
        "\n",
        "        # Set vocabulary size, maximum length, and embedding dimension\n",
        "        self.vocab_size = vocab_size\n",
        "        self.max_len = max_len\n",
        "        self.embedding_dim = embedding_dim\n",
        "\n",
        "        # Token embeddings layer\n",
        "        self.token_embeddings = nn.Embedding(self.vocab_size, self.embedding_dim)\n",
        "\n",
        "        # Position embeddings layer with precomputed weights and frozen\n",
        "        weight_position_embeddings = self.get_position_encoding(max_len, embedding_dim)\n",
        "        self.position_embeddings = nn.Embedding(self.max_len, self.embedding_dim, _weight=weight_position_embeddings, _freeze=True)\n",
        "\n",
        "    def get_position_encoding(self, seq_length, hidden_size, n=10000):\n",
        "        # Generate position encoding based on the given sequence length, hidden size, and frequency (n)\n",
        "\n",
        "        position_enc = torch.zeros(seq_length, hidden_size)\n",
        "\n",
        "        for pos in range(seq_length):\n",
        "            for i in range(hidden_size // 2):\n",
        "                position_enc[pos, 2*i] = torch.sin(torch.tensor(pos / (n**(2*i / hidden_size))))\n",
        "                position_enc[pos, 2*i+1] = torch.cos(torch.tensor(pos / (n**((2*i+1) / hidden_size))))\n",
        "\n",
        "        return position_enc\n",
        "\n",
        "    def forward(self, inputs):\n",
        "        # Forward pass of the TokenPositionEmbeddings module\n",
        "\n",
        "        # inputs is of shape [batch_size, max_len]\n",
        "        batch_size, max_len = inputs.shape\n",
        "\n",
        "        # Create position indices\n",
        "        position_indices = torch.arange(max_len).to(device)\n",
        "\n",
        "        # Get token embeddings\n",
        "        token_embeddings = self.token_embeddings(inputs)\n",
        "\n",
        "        # Get position embeddings\n",
        "        position_embeddings = self.position_embeddings(position_indices)\n",
        "\n",
        "        # Add both token and position embeddings\n",
        "        embeddings = token_embeddings + position_embeddings\n",
        "\n",
        "        return embeddings\n"
      ],
      "metadata": {
        "id": "PSBLJfAYj8OR"
      },
      "execution_count": 67,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## <font color='red'> ScaledDotProduct and Multihead attention (the secret sauce) </font>"
      ],
      "metadata": {
        "id": "x1pT5MeljYG7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn as nn\n",
        "\n",
        "class ScaledDotProductAttention(nn.Module):\n",
        "    def __init__(self):\n",
        "        # Initialize ScaledDotProductAttention module\n",
        "        super().__init__()\n",
        "\n",
        "    def forward(self, query, key, value, mask=None):\n",
        "        # Forward pass of the ScaledDotProductAttention module\n",
        "\n",
        "        # query, key, and value shapes: [B, nh, T, hs]\n",
        "        B, nh, T, hs = query.shape\n",
        "\n",
        "        # Transpose key for matrix multiplication\n",
        "        key = key.transpose(-2, -1)\n",
        "\n",
        "        # Calculate attention scores\n",
        "        attention_score = torch.matmul(query, key)\n",
        "\n",
        "        # Scale the attention scores\n",
        "        attention_score = attention_score / torch.sqrt(torch.tensor(hs))\n",
        "\n",
        "        # Apply mask if provided\n",
        "        if mask is not None:\n",
        "            attention_score += (mask * -1e9)\n",
        "\n",
        "        # Apply softmax to get attention weights\n",
        "        attention_weights = torch.softmax(attention_score, dim=-1)\n",
        "\n",
        "        # Calculate output using attention weights and value\n",
        "        output = torch.matmul(attention_weights, value)\n",
        "        # Output shape: [B, nh, T, hs]\n",
        "\n",
        "        return output\n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, num_heads, embed_size, key_dim, query_dim, value_dim, mask=False):\n",
        "        # Initialize MultiHeadAttention module with specified parameters\n",
        "        super().__init__()\n",
        "\n",
        "        self.num_heads = num_heads\n",
        "        self.head_size = embed_size // num_heads\n",
        "        self.key_dim = key_dim\n",
        "        self.query_dim = query_dim\n",
        "        self.value_dim = value_dim\n",
        "        self.mask = mask\n",
        "\n",
        "        # ScaledDotProductAttention module\n",
        "        self.attention = ScaledDotProductAttention()\n",
        "\n",
        "        # Linear layers for key, query, value, and final layer\n",
        "        self.key_layer = nn.Linear(embed_size, key_dim).to(device)\n",
        "        self.query_layer = nn.Linear(embed_size, query_dim).to(device)\n",
        "        self.value_layer = nn.Linear(embed_size, value_dim).to(device)\n",
        "        self.final_layer = nn.Linear(value_dim, embed_size).to(device)\n",
        "\n",
        "    def forward(self, query, key, value, mask=None):\n",
        "        # Forward pass of the MultiHeadAttention module\n",
        "\n",
        "        # Input shape: [B, T, C]\n",
        "        B, T, C = query.shape\n",
        "\n",
        "        # Apply linear layers to key, query, and value\n",
        "        query = self.query_layer(query)\n",
        "        key = self.key_layer(key)\n",
        "        value = self.value_layer(value)\n",
        "\n",
        "        # Reshape key, query, and value for multiple heads\n",
        "        query = query.view(B, T, self.num_heads, -1).transpose(1, 2)\n",
        "        key = key.view(B, T, self.num_heads, -1).transpose(1, 2)\n",
        "        value = value.view(B, T, self.num_heads, -1).transpose(1, 2)\n",
        "\n",
        "        # Apply ScaledDotProductAttention\n",
        "        scaled_attention = self.attention(query, key, value, mask)\n",
        "        # Scaled_attention shape: [B, T, nh, hs]\n",
        "\n",
        "        # Transpose and reshape scaled_attention\n",
        "        scaled_attention = scaled_attention.transpose(1, 2).contiguous().view(B, T, -1)\n",
        "        # Scaled_attention shape: [B, T, nh*hs]\n",
        "\n",
        "        # Apply final linear layer\n",
        "        output = self.final_layer(scaled_attention)\n",
        "        # Output shape: [B, T, C]\n",
        "\n",
        "        return output\n"
      ],
      "metadata": {
        "id": "D-gOpTQ4iLkT"
      },
      "execution_count": 68,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## <font color='red'> The Encoder block </font>"
      ],
      "metadata": {
        "id": "2GX2AtPmkLA9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn as nn\n",
        "\n",
        "class EncoderBlock(nn.Module):\n",
        "    def __init__(self, num_heads, embed_size, key_dim, query_dim, value_dim):\n",
        "        # Initialize EncoderBlock module with specified parameters\n",
        "        super().__init__()\n",
        "\n",
        "        # MultiHeadAttention module\n",
        "        self.MultiHeadAttention = MultiHeadAttention(num_heads, embed_size, key_dim, query_dim, value_dim)\n",
        "\n",
        "        # Layer normalization\n",
        "        self.layer_norm = nn.LayerNorm(embed_size).to(device)\n",
        "\n",
        "        # Feed-forward neural network\n",
        "        self.feed_forward = nn.Sequential(\n",
        "            nn.Linear(embed_size, 4 * embed_size),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(4 * embed_size, embed_size)\n",
        "        ).to(device)\n",
        "\n",
        "    def forward(self, inputs, mask):\n",
        "        # Forward pass of the EncoderBlock module\n",
        "\n",
        "        # Apply MultiHeadAttention\n",
        "        attention = self.MultiHeadAttention(inputs, inputs, inputs, mask)\n",
        "\n",
        "        # Apply layer normalization and residual connection\n",
        "        normalized_op1 = self.layer_norm(inputs + attention)\n",
        "\n",
        "        # Apply feed-forward neural network\n",
        "        feed_forward_op = self.feed_forward(normalized_op1)\n",
        "\n",
        "        # Apply layer normalization and another residual connection\n",
        "        normalized_op2 = self.layer_norm(normalized_op1 + feed_forward_op)\n",
        "\n",
        "        return normalized_op2\n",
        "\n",
        "\n",
        "class Encoder(nn.Module):\n",
        "    def __init__(self, src_max_length, embedding_dim, key_dim, query_dim, value_dim, src_vocab_size, dropout_rate, num_blocks, num_heads):\n",
        "        # Initialize Encoder module with specified parameters\n",
        "        super().__init__()\n",
        "\n",
        "        # Maximum length for positional embeddings\n",
        "        self.max_length = src_max_length\n",
        "\n",
        "        # TokenPositionEmbeddings module\n",
        "        self.token_position_embeddings = TokenPositionEmbeddings(src_vocab_size, src_max_length, embedding_dim)\n",
        "\n",
        "        # List of EncoderBlocks\n",
        "        self.encoder_stack = [EncoderBlock(num_heads, embedding_dim, key_dim, query_dim, value_dim) for _ in range(num_blocks)]\n",
        "\n",
        "    def forward(self, inputs, mask):\n",
        "        # Forward pass of the Encoder module\n",
        "\n",
        "        # Apply token and position embeddings\n",
        "        x = self.token_position_embeddings(inputs)\n",
        "\n",
        "        # Iterate through the encoder stack\n",
        "        for encoder_block in self.encoder_stack:\n",
        "            x = encoder_block(x, mask)\n",
        "\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "s8v-ic5Cj3SF"
      },
      "execution_count": 69,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## <font color='red'> The Decoder Block </font>"
      ],
      "metadata": {
        "id": "hE-RmRMekNTH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn as nn\n",
        "\n",
        "class DecoderBlock(nn.Module):\n",
        "    def __init__(self, num_heads, embed_size, key_dim, query_dim, value_dim):\n",
        "        # Initialize DecoderBlock module with specified parameters\n",
        "        super().__init__()\n",
        "\n",
        "        # MultiHeadAttention module for masked attention\n",
        "        self.maskedMultiHeadAttention = MultiHeadAttention(num_heads, embed_size, key_dim, query_dim, value_dim)\n",
        "\n",
        "        # MultiHeadAttention module for attention with encoder output\n",
        "        self.multihead_attention = MultiHeadAttention(num_heads, embed_size, key_dim, query_dim, value_dim)\n",
        "\n",
        "        # Layer normalization\n",
        "        self.layer_norm = nn.LayerNorm(embed_size).to(device)\n",
        "\n",
        "        # Feed-forward neural network\n",
        "        self.feed_forward = nn.Sequential(\n",
        "            nn.Linear(embed_size, embed_size * 4),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(embed_size * 4, embed_size)\n",
        "        ).to(device)\n",
        "\n",
        "    def forward(self, dec_input, enc_output, lookahead_mask, padding_mask):\n",
        "        # Forward pass of the DecoderBlock module\n",
        "\n",
        "        # Apply masked multi-head attention\n",
        "        masked_attn_output = self.maskedMultiHeadAttention(dec_input, dec_input, dec_input, lookahead_mask)\n",
        "\n",
        "        # Apply layer normalization and residual connection\n",
        "        normalized_op1 = self.layer_norm(dec_input + masked_attn_output)\n",
        "\n",
        "        # Apply multi-head attention with encoder output\n",
        "        attn_output = self.multihead_attention(normalized_op1, enc_output, enc_output, padding_mask)\n",
        "\n",
        "        # Apply layer normalization and another residual connection\n",
        "        normalized_op2 = self.layer_norm(normalized_op1 + attn_output)\n",
        "\n",
        "        # Apply feed-forward neural network\n",
        "        feed_forward_output = self.feed_forward(normalized_op2)\n",
        "\n",
        "        # Apply layer normalization and another residual connection\n",
        "        normalized_op3 = self.layer_norm(normalized_op2 + feed_forward_output)\n",
        "\n",
        "        return normalized_op3\n",
        "\n",
        "class Decoder(nn.Module):\n",
        "    def __init__(self, tar_max_length, embedding_dim, key_dim, value_dim, query_dim, tar_vocab_size, dropout_rate, num_blocks, num_heads, device='cpu'):\n",
        "        # Initialize Decoder module with specified parameters\n",
        "        super().__init__()\n",
        "\n",
        "        # Maximum length for positional embeddings\n",
        "        self.max_length = tar_max_length\n",
        "\n",
        "        # TokenPositionEmbeddings module\n",
        "        self.token_position_embeddings = TokenPositionEmbeddings(tar_vocab_size, tar_max_length, embedding_dim)\n",
        "\n",
        "        # List of DecoderBlocks\n",
        "        self.decoder_stack = [DecoderBlock(num_heads, embedding_dim, key_dim, value_dim, query_dim) for _ in range(num_blocks)]\n",
        "\n",
        "        self.device = device\n",
        "\n",
        "    def forward(self, inputs, enc_output, lookahead_mask, padding_mask):\n",
        "        # Forward pass of the Decoder module\n",
        "\n",
        "        # Apply token and position embeddings\n",
        "        x = self.token_position_embeddings(inputs)\n",
        "\n",
        "        # Iterate through the decoder stack\n",
        "        for decoder_block in self.decoder_stack:\n",
        "            x = decoder_block(x, enc_output, lookahead_mask, padding_mask)\n",
        "\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "So8RbM5FkNrd"
      },
      "execution_count": 70,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## <font color='red'> The complete Transformer Model </font>"
      ],
      "metadata": {
        "id": "EqjZAiTAkRgI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn as nn\n",
        "\n",
        "class TransformerModel(nn.Module):\n",
        "    def __init__(self, src_max_length, tar_max_length, embedding_dim, key_dim, query_dim, value_dim, src_vocab_size, tar_vocab_size, dropout_rate, num_blocks, num_heads, device='cpu'):\n",
        "        # Initialize TransformerModel module with specified parameters\n",
        "        super().__init__()\n",
        "\n",
        "        # Encoder module\n",
        "        self.encoder = Encoder(src_max_length, embedding_dim, key_dim, query_dim, value_dim, src_vocab_size, dropout_rate, num_blocks, num_heads)\n",
        "\n",
        "        # Decoder module\n",
        "        self.decoder = Decoder(tar_max_length, embedding_dim, key_dim, query_dim, value_dim, tar_vocab_size, dropout_rate, num_blocks, num_heads)\n",
        "\n",
        "        # Final linear layer\n",
        "        self.final_layer = nn.Linear(embedding_dim, tar_vocab_size)\n",
        "\n",
        "    def create_padding_mask(self, inputs):\n",
        "        # Create padding mask for inputs\n",
        "        mask = torch.zeros(inputs.shape[0], inputs.shape[1]).to(device)\n",
        "        mask = mask.masked_fill(inputs == 0, 1)\n",
        "        mask = mask.view(inputs.shape[0], 1, 1, inputs.shape[1])\n",
        "        return mask\n",
        "\n",
        "    def create_lookahead_mask(self, inputs):\n",
        "        # Create lookahead mask for inputs\n",
        "        mask = torch.triu(torch.ones((inputs.shape[1], inputs.shape[1])), diagonal=1)\n",
        "        return mask\n",
        "\n",
        "    def forward(self, enc_inputs, dec_inputs, target):\n",
        "        # Forward pass of the TransformerModel module\n",
        "\n",
        "        # Create padding mask for encoder inputs\n",
        "        padding_mask_enc = self.create_padding_mask(enc_inputs)\n",
        "\n",
        "        # Create padding mask and lookahead mask for decoder inputs\n",
        "        padding_mask_dec = self.create_padding_mask(dec_inputs)\n",
        "        lookahead_mask_dec = self.create_lookahead_mask(dec_inputs).to(device)\n",
        "\n",
        "        # Combine padding mask and lookahead mask for decoder\n",
        "        dec_mask = torch.max(padding_mask_dec, lookahead_mask_dec)\n",
        "\n",
        "        # Forward pass through encoder\n",
        "        enc_output = self.encoder(enc_inputs, padding_mask_enc)\n",
        "\n",
        "        # Forward pass through decoder\n",
        "        dec_output = self.decoder(dec_inputs, enc_output, dec_mask, padding_mask_dec)\n",
        "\n",
        "        # Forward pass through final linear layer\n",
        "        output = self.final_layer(dec_output)\n",
        "\n",
        "        loss = None\n",
        "        if target is not None:\n",
        "            # Calculate CrossEntropyLoss if target is provided\n",
        "            loss_fct = nn.CrossEntropyLoss(ignore_index=0)\n",
        "            output = output.reshape(-1, output.shape[-1])\n",
        "            target = target.reshape(-1)\n",
        "            loss = loss_fct(output, target)\n",
        "\n",
        "        return output, loss\n"
      ],
      "metadata": {
        "id": "BKtmoMjlkRwJ"
      },
      "execution_count": 71,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## <font color='red'> Model Training </font>"
      ],
      "metadata": {
        "id": "4jUSwweAkqLs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch.optim import Adam\n",
        "from torch.nn import CrossEntropyLoss\n",
        "from torch.optim.lr_scheduler import LambdaLR\n",
        "\n",
        "def get_batch_data(batch_size, is_train=True):\n",
        "    # Get batch data for training or testing\n",
        "    if is_train:\n",
        "        eng = train_eng\n",
        "        ger = train_ger\n",
        "    else:\n",
        "        eng = test_eng\n",
        "        ger = test_ger\n",
        "\n",
        "    # Get random index as batch starting index\n",
        "    batch_start = np.random.randint(0, len(eng) - batch_size)\n",
        "    batch_end_index = batch_start + batch_size\n",
        "    batch_eng = eng[batch_start:batch_end_index]\n",
        "    batch_ger = ger[batch_start:batch_end_index]\n",
        "\n",
        "    return batch_eng, batch_ger\n",
        "\n",
        "# Optimizer\n",
        "if __name__ == '__main__':\n",
        "    ger_max_len = 33\n",
        "    eng_max_len = 33\n",
        "    ger_vocab_size = 8663\n",
        "    eng_vocab_size = 5627\n",
        "    embedding_size = 256\n",
        "    num_blocks = 6\n",
        "    num_heads = 8\n",
        "    key_dim = 64\n",
        "    query_dim = 64\n",
        "    value_dim = 64\n",
        "    epochs =2000\n",
        "    batch_size = 32\n",
        "    eval_iters = 100\n",
        "\n",
        "    # Initialize DataLoader\n",
        "    dataset = TranslationDataloader('/content/drive/MyDrive/kData/english-german_60k.pkl', 15000)\n",
        "    train_eng, train_ger, test_eng, test_ger = dataset.get_encoded_data()\n",
        "\n",
        "    # Initialize Transformer Model\n",
        "    model = TransformerModel(src_max_length=eng_max_len, tar_max_length=ger_max_len, embedding_dim=embedding_size, key_dim=key_dim, query_dim=query_dim,\n",
        "                             value_dim=value_dim, src_vocab_size=eng_vocab_size, tar_vocab_size=ger_vocab_size, dropout_rate=0.1, num_blocks=num_blocks,\n",
        "                             num_heads=num_heads, device=device)\n",
        "\n",
        "    # Move model to device\n",
        "    model = model.to(device)\n",
        "\n",
        "    # Learning rate scheduling function\n",
        "    def rate(step, model_size, factor, warmup):\n",
        "        if step == 0:\n",
        "            step = 1\n",
        "        return factor * (\n",
        "            model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))\n",
        "        )\n",
        "\n",
        "    # Optimizer and learning rate scheduler\n",
        "    optimizer = Adam(\n",
        "        model.parameters(), lr=1.0, betas=(0.9, 0.98), eps=1e-9)\n",
        "\n",
        "    lr_scheduler = LambdaLR(\n",
        "        optimizer=optimizer,\n",
        "        lr_lambda=lambda step: rate(\n",
        "            step, embedding_size, factor=1, warmup=1000\n",
        "        ),\n",
        "    )\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def estimate_loss():\n",
        "        # Estimate loss on training and test splits on random 100 batches\n",
        "        out = []\n",
        "        model.eval()\n",
        "\n",
        "        for split in [True, False]:\n",
        "            losses = torch.zeros(eval_iters)\n",
        "            for k in range(eval_iters):\n",
        "                eng_batch, ger_batch = get_batch_data(batch_size, split)\n",
        "                encoder_input = eng_batch[:, 1:].to(device)\n",
        "                decoder_input = ger_batch[:, :-1].to(device)\n",
        "                decoder_target = ger_batch[:, 1:].to(device)\n",
        "                logits, loss = model(encoder_input, decoder_input, decoder_target)\n",
        "                losses[k] = loss.item()\n",
        "            out.append(losses.mean())\n",
        "        model.train()\n",
        "        return out\n",
        "    min_loss=1000\n",
        "    for epoch in range(epochs):\n",
        "        eng_batch, ger_batch = get_batch_data(batch_size, True)\n",
        "        encoder_input = eng_batch[:, 1:].to(device)\n",
        "        decoder_input = ger_batch[:, :-1].to(device)\n",
        "        decoder_target = ger_batch[:, 1:].to(device)\n",
        "        logits, loss = model(encoder_input, decoder_input, decoder_target)\n",
        "\n",
        "        if epoch % 200 == 0:\n",
        "            train_loss, test_loss = estimate_loss()\n",
        "            print('Epoch: ', epoch, 'Train Loss: ', train_loss.item(), 'Test Loss: ', test_loss.item())\n",
        "            if(train_loss<min_loss):\n",
        "              min_loss=train_loss\n",
        "              torch.save(model.state_dict(),'transformer_model.pt')\n",
        "        optimizer.zero_grad(set_to_none=True)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        lr_scheduler.step()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QFfVeiwHk7aA",
        "outputId": "ce48fbe3-76d6-4c18-cd6d-6a0c38c55494"
      },
      "execution_count": 72,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "12750\n",
            "\n",
            "Train Data Size:  12750\n",
            "Test Data Size:  2250\n",
            "Max German Sentence Length:  33\n",
            "Max English Sentence Length:  33\n",
            "English Vocab Size:  5627\n",
            "German Vocab Size:  8663\n",
            "Epoch:  0 Train Loss:  9.216261863708496 Test Loss:  9.209721565246582\n",
            "Epoch:  200 Train Loss:  6.0287275314331055 Test Loss:  5.857973575592041\n",
            "Epoch:  400 Train Loss:  5.486483573913574 Test Loss:  5.339689254760742\n",
            "Epoch:  600 Train Loss:  4.954954624176025 Test Loss:  4.979308605194092\n",
            "Epoch:  800 Train Loss:  4.376029014587402 Test Loss:  4.67732048034668\n",
            "Epoch:  1000 Train Loss:  3.8992648124694824 Test Loss:  4.501443862915039\n",
            "Epoch:  1200 Train Loss:  3.683950424194336 Test Loss:  4.332326889038086\n",
            "Epoch:  1400 Train Loss:  3.3603010177612305 Test Loss:  4.217817306518555\n",
            "Epoch:  1600 Train Loss:  3.1712257862091064 Test Loss:  4.190222263336182\n",
            "Epoch:  1800 Train Loss:  2.857576370239258 Test Loss:  4.1197662353515625\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "JuPnauhxk8nG"
      },
      "execution_count": 72,
      "outputs": []
    }
  ]
}