attention_is_all_you_need.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class InputEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model:int, vocab_size:int):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.vocab_size = vocab_size\n",
    "        self.embedding = nn.Embedding(vocab_size, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.embedding(x)*math.sqrt(self.d_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositonalEncoding(nn.Module):\n",
    "    def __init__(self, d_model:int, seq_len:int, dropout:float) -> None:\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.seq_len = seq_len\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        #create a matrix of shape (seq_len, d_model)\n",
    "        pe = torch.zeros(seq_len, d_model)\n",
    "        #create a vector of shape(seq_len, 1)\n",
    "        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "        #apply the sin to even positions\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        #apply the cos to odd positions\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0) #(1, seq_len, d_model)  # adding the batch dim\n",
    "        self.register_buffer('pe', pe) # registering the positional encoding as a buffer; it means we will store these tensors when we save the model Add a buffer to the module.\n",
    "        #This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm's running_mean is not a parameter, but is part of the module's state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module's state_dict.\n",
    "        #Buffers can be accessed as attributes using given names.\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)\n",
    "        return self.dropout(x)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNormalization(nn.Module):\n",
    "    def __init__(self, eps: float=1e-6):\n",
    "        super().__init__()\n",
    "        self.eps = eps\n",
    "        self.alpha = nn.Parameter(torch.ones(1))\n",
    "        self.bias = nn.Parameter(torch.zeros(1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        return self.alpha * (x - mean) / (std + self.eps) + self.bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForwardBlock(nn.Module):\n",
    "    def __init__(self, d_model:int, d_ff:int, dropout:float):\n",
    "        super().__init__()\n",
    "        self.linear1 = nn.Linear(d_model, d_ff)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear2 = nn.Linear(d_ff, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #(batch, seq_len, d_model) --> (batch, seq_len, d_ff) --> (batch, seq_len, d_model)\n",
    "        return self.linear2(self.dropout(torch.relu(self.linear1(x))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttentionBlock(nn.Module):\n",
    "    def __init__(self, d_model:int, h:int, dropout:float) -> None:\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.h = h\n",
    "        assert d_model%h  == 0, \"d_model is not divisible by h\"\n",
    "\n",
    "        self.d_k = d_model//h\n",
    "        self.w_q = nn.Linear(d_model, d_model) #Wq\n",
    "        self.w_k = nn.Linear(d_model, d_model) #Wk\n",
    "        self.w_v = nn.Linear(d_model, d_model) #Wv\n",
    "\n",
    "        self.w_o = nn.Linear(d_model, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "\n",
    "    @staticmethod\n",
    "    def attention(query, key, value, mask, dropout:nn.Dropout):\n",
    "        d_k = query.shape[-1]\n",
    "        #(batch, h, seq_len, d_k) --> (batch, seq_len, d_model)\n",
    "        attention_scores = (query @ key.transpose(-2, -1))/math.sqrt(d_k) #<-- needs visualization\n",
    "        if mask is not None:\n",
    "            attention_scores.masked_fill_(mask == 0, -1e9)\n",
    "        attention_scores = attention_scores.softmax(dim=-1) #batch , h, seq_len, seq_len\n",
    "        if dropout is not None:\n",
    "            attention_scores = dropout(attention_scores)\n",
    "        \n",
    "        return (attention_scores @ value), attention_scores\n",
    "\n",
    "\n",
    "    def forward(self, q, k, v, mask):\n",
    "        query = self.w_q(q)\n",
    "        key = self.w_k(k)\n",
    "        value = self.w_v(v)\n",
    "\n",
    "        # (batch, seq_len, d_model) --> (batch, h, seq_len, d_k) // we want each head to look over the seq_len entirely but a small part of the embedding\n",
    "        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1,2)\n",
    "        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1,2)\n",
    "        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(1,2)\n",
    "\n",
    "        x, self.attention_scores = MultiHeadAttentionBlock.attention(query, key, value, mask, self.dropout)\n",
    "\n",
    "        x = x.transpose(1,2).contigous().view(x.shape[0], -1, self.h * self.d_k) #<-- check what is supposed to happen here\n",
    "\n",
    "        #(batch, seq_len, d_model) --> same \n",
    "        return self.w_o(x) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResidualConnection(nn.Module):\n",
    "    def __init__(self, dropout:float):\n",
    "        super().__init__()\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, x, sublayer):\n",
    "        return x + self.dropout(sublayer(self.norm(x))) # in this paper the reverse was done, but idk stability demands sublayer(norm())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    def __init__(self, self_attention_block:MultiHeadAttentionBlock, feed_forward_block:FeedForwardBlock, dropout: float) -> None:\n",
    "        super().__init__()\n",
    "        self.self_attention_block = self_attention_block\n",
    "        self.feed_forward_block = feed_forward_block\n",
    "        self.residual_connections = nn.ModuleList([ResidualConnection(dropout) for _ in range(2)])\n",
    "\n",
    "    def forward(self, x, src_mask):\n",
    "        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, layers:nn.ModuleList):\n",
    "        super().__init__()\n",
    "        self.layers = layers\n",
    "        self.norm = LayerNormalization()\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, mask)\n",
    "        return self.norm(x)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderBlock(nn.Module):\n",
    "    def __init__(self, self_attention_block:MultiHeadAttentionBlock, cross_attention_block: MultiHeadAttentionBlock, feed_forward_block: FeedForwardBlock, dropout: float):\n",
    "        super().__init__()\n",
    "        self.self_attention_block = self_attention_block\n",
    "        self.cross_attention_block = cross_attention_block\n",
    "        self.feed_forward_block = feed_forward_block\n",
    "        self.residual_connections = nn.ModuleList([ResidualConnection(dropout) for _ in range(3)])\n",
    "        \n",
    "    def forward(self, x, encoder_output, src_mask, tgt_mask):\n",
    "        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))\n",
    "        x = self.residual_connections[1](x, lambda x: self.self_attention_block(x, encoder_output, encoder_output, src_mask))\n",
    "        x = self.residual_connections[2](x, lambda x: self.feed_forward_block(x, self.feed_forward_block))\n",
    "        return x\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, layers:nn.ModuleList) -> None:\n",
    "        super().__init__()\n",
    "        self.layers = layers\n",
    "        self.norm = LayerNormalization()\n",
    "\n",
    "    def forward(self, x, encoder_output, src_mask, tgt_mask):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, encoder_output, src_mask, tgt_mask)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProjectionLayer(nn.Module):\n",
    "    def __init__(self, d_model:int, vocab_size:int):\n",
    "        super().__init__()\n",
    "        self.proj = nn.Linear(d_model, vocab_size)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # (batch, seq_len, d_model) -> (batch, seq_len, vocab_size)\n",
    "        return torch.log_softmax(self.proj(x), dim = -1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: InputEmbeddings, tgt_embed: InputEmbeddings, src_pos: PositonalEncoding, tgt_pos:PositonalEncoding, projection_layer: ProjectionLayer) -> None:\n",
    "        super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_embed = src_embed\n",
    "        self.src_pos = src_pos\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.tgt_pos = tgt_pos\n",
    "        self.projection_layer = projection_layer\n",
    "\n",
    "    def encode(self, src, src_mask):\n",
    "        src = self.src_embed(src)\n",
    "        src = self.src_pos(src)\n",
    "        return self.encoder(src, src_mask)\n",
    "    \n",
    "    def decode(self, encoder_output, src_mask, tgt, tgt_mask):\n",
    "        tgt = self.tgt_embed(tgt)\n",
    "        tgt = self.tgt_pos(tgt)\n",
    "        return self.decoder(tgt, encoder_output, src_mask, tgt_mask)\n",
    "    \n",
    "    def project(self, x):\n",
    "        return self.projection_layer(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int, tgt_seq_len: int, d_model: 512, N: int = 6, h: int = 8, dropout: float = 0.1, d_ff: int = 2048) -> Transformer:\n",
    "    #Create the embedding layer\n",
    "    src_embed = InputEmbeddings(d_model, src_vocab_size)\n",
    "    tgt_embed = InputEmbeddings(d_model, tgt_vocab_size)\n",
    "\n",
    "    #create the positional encoding layers\n",
    "    src_pos = PositionalEncoding(d_model, src_seq_len, dropout)\n",
    "    tgt_pos = PositionalEncoding(d_model, tgt_seq_len, dropout)\n",
    "\n",
    "    #create encoder blocks\n",
    "    encoder_blocks = []\n",
    "    for _ in range(N):\n",
    "        encoder_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)\n",
    "        feed_forward_block = FeedForwardBlock(d_model, h, dropout)\n",
    "        encoder_block = EncoderBlock(encoder_self_attention_block, feed_forward_block, dropout)\n",
    "        encoder_blocks.append(encoder_block)\n",
    "    \n",
    "    #create the decoder blocks\n",
    "    decoder_blocks = []\n",
    "    for _ in range(N):\n",
    "        decoder_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)\n",
    "        decoder_cross_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)\n",
    "        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)\n",
    "        decoder_block = DecoderBlock(decoder_self_attention_block, decoder_cross_attention_block, feed_forward_block, dropout)\n",
    "        decoder_blocks.append(decoder_block)\n",
    "    \n",
    "    #create the encoder and decoder\n",
    "    encoder = Encoder(nn.ModuleList(encoder_blocks))\n",
    "    decoder = Decoder(nn.ModuleList(decoder_blocks))\n",
    "\n",
    "    #create a projection layer\n",
    "    projection_layer = ProjectionLayer(d_model, tgt_vocab_size)\n",
    "    \n",
    "    #create the transformer\n",
    "    transformer = Transformer(encoder, decoder, src_embed, tgt_embed, src_pos, tgt_pos, projection_layer)\n",
    "\n",
    "    #initialize the transformer\n",
    "    for p in transformer.parameters():\n",
    "        if p.dim()>1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "    \n",
    "    return transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Arc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from datasets import "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.090001"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eps = 10**-6\n",
    "eps + 1.09"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 4.,  5.,  6.,  7.],\n",
       "         [ 9., 10., 11., 12.],\n",
       "         [ 6.,  7.,  8.,  9.]]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = torch.Tensor([[1, 2, 3, 4], [5, 6, 7, 8, ], [1, 2, 3, 4]]).unsqueeze(0)\n",
    "b = torch.Tensor([[3, 3, 3, 3],[4, 4, 4, 4], [5, 5, 5, 5]]).unsqueeze(0)\n",
    "\n",
    "a = a + b[:, :b.shape[1], : ]\n",
    "a\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 5.5000],\n",
      "         [10.5000],\n",
      "         [ 7.5000]]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 1])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = a.mean(dim = -1, keepdim=True)#.unsqueeze(1)\n",
    "print(m)\n",
    "m.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 1])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[2.],\n",
       "        [3.],\n",
       "        [4.]], dtype=torch.float64)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test = torch.arange(2, 5, dtype=torch.float64).unsqueeze(1)\n",
    "s = test.shape\n",
    "print(s)\n",
    "test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([2., 3., 4.], dtype=torch.float64)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test.squeeze(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1.0000e+00, 2.7183e+00, 7.3891e+00, 2.0086e+01, 5.4598e+01, 1.4841e+02,\n",
       "        4.0343e+02, 1.0966e+03, 2.9810e+03, 8.1031e+03], dtype=torch.float64)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "b  = torch.arange(0, 10, dtype=torch.float64)\n",
    "\n",
    "a = torch.exp(b)\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "print(b)\n",
    "b = b.view(5, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 1.],\n",
       "        [2., 3.],\n",
       "        [4., 5.],\n",
       "        [6., 7.],\n",
       "        [8., 9.]], dtype=torch.float64)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}