architecture_details.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Architecture of Transformer\n",
    "\n",
    "![](img/attention_arch.jpeg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input Embedding\n",
    "\n",
    "- It takes the input and convert it into an embedding\n",
    "- It allows to convert the original sentence into a vector (ex 512 dimensions)\n",
    "- **Note**: \n",
    "   - `d_model` is the model's dimension\n",
    "   - `vocab_size` is the number of words in the vocabulary\n",
    "\n",
    "![](img/input_embeded.png)\n",
    "\n",
    "$$\n",
    "\\text{Embedding weights} \\times \\sqrt{d_{\\text{model}}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class InputEmbeddings(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 = Embedding(vocab_size, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.embedding(x) * math.sqrt(self.d_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Positional Encoding\n",
    "\n",
    "- Convey to the model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(Module):\n",
    "    def __init__(self, d_model: int, seq_len: int, dropout: float):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.seq_len = seq_len\n",
    "        self.dropout = Dropout(dropout)\n",
    "\n",
    "        # Create a matrix of shape (seqlen, d_model)\n",
    "        pe = torch.zeros(seq_len, d_model)\n",
    "        # Create a vector of shape(seqlen, 1)\n",
    "        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)\n",
    "\n",
    "        div_term = torch.exp(\n",
    "            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)\n",
    "        )\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",
    "\n",
    "        pe = pe.unsqueeze(0)  # (1, seq_len, d_model)\n",
    "\n",
    "        self.register_buffer(\"pe\", pe)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x += (self.pe[:, : x.shape[1], :]).requires_grad_(False)\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNormalization(Module):\n",
    "    def __init__(self, eps: float = 10**6):\n",
    "        super().__init__()\n",
    "        self.eps = eps\n",
    "        self.alpha = Parameter(torch.ones(1))  # Multiplied\n",
    "        self.bias = Parameter(torch.zeros(1))  # Added\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(dim=-1, keepdim=True)\n",
    "        std = x.std(dim=-1, keepdim=True)\n",
    "        return self.alpha * (x - mean) / (std + self.eps) + self.bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForwardBlock(Module):\n",
    "    def __init__(self, d_model: int, d_ff: int, dropout: float):\n",
    "        super().__init__()\n",
    "        self.linear_1 = Linear(d_model, d_ff)  # W1 and b1\n",
    "        self.dropout = Dropout(dropout)\n",
    "        self.linear_2 = Linear(d_ff, d_model)  # W2 and B2\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.linear_2(self.dropout(torch.relu(self.linear_1(x))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttentionBlock(Module):\n",
    "    def __init__(self, d_model: int, h: int, dropout: float):\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 = Linear(d_model, d_model)  # Wq\n",
    "        self.w_k = Linear(d_model, d_model)  # Wk\n",
    "        self.w_v = Linear(d_model, d_model)  # Wv\n",
    "\n",
    "        self.w_o = Linear(d_model, d_model)  # Wo\n",
    "        self.dropout = Dropout(dropout)\n",
    "\n",
    "    @staticmethod\n",
    "    def attention(query, key, value, mask, dropout: Dropout):\n",
    "        d_k = query.shape[-1]\n",
    "\n",
    "        # (Batch, h, Seq_len, d_k) --> (Batch, h, Seq_Len, Seq_Len)\n",
    "        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)\n",
    "        if mask is not None:\n",
    "            attention_scores.masked_fill_(mask == 0, -1e9)\n",
    "        attention_scores = attention_scores.softmax(\n",
    "            dim=-1\n",
    "        )  # (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",
    "    def forward(self, q, k, v, mask):\n",
    "        query = self.w_q(q)  # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, d_model)\n",
    "        key = self.w_k(k)  # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, d_model)\n",
    "        value = self.w_v(v)  # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, d_model)\n",
    "\n",
    "        # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, h, d_k) --> (Batch, h, Seq_Len, d_k)\n",
    "        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(\n",
    "            1, 2\n",
    "        )\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(\n",
    "            1, 2\n",
    "        )\n",
    "\n",
    "        x, self.attention_scores = MultiHeadAttentionBlock.attention(\n",
    "            query, key, value, mask, self.dropout\n",
    "        )\n",
    "\n",
    "        # (Batch, h, seq_len, d_k) --> (Batch, h, seq_len, d_k) --> (Batch, Seq_Len, d_model)\n",
    "        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)\n",
    "\n",
    "        # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, d_model)\n",
    "        return self.w_o(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResidualConnection(Module):\n",
    "\n",
    "    def __init__(self, dropout: float):\n",
    "        super().__init__()\n",
    "        self.dropout = Dropout(dropout)\n",
    "        self.norm = LayerNormalization()\n",
    "\n",
    "    def forward(self, x, sublayer):\n",
    "        return x + self.dropout(sublayer(self.norm(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        self_attention_block: MultiHeadAttentionBlock,\n",
    "        feed_forward_block: FeedForwardBlock,\n",
    "        dropout: float,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.self_attention_block = self_attention_block\n",
    "        self.feed_forward_block = feed_forward_block\n",
    "        self.residual_connections = ModuleList(\n",
    "            [ResidualConnection(dropout) for _ in range(2)]\n",
    "        )\n",
    "\n",
    "    def forward(self, x, src_mask):\n",
    "        x = self.residual_connections[0](\n",
    "            x, lambda x: self.self_attention_block(x, x, x, src_mask)\n",
    "        )\n",
    "        x = self.residual_connections[1](x, self.feed_forward_block)\n",
    "        return x\n",
    "\n",
    "\n",
    "class Encoder(Module):\n",
    "    def __init__(self, layers: 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderBlock(Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        features: int,\n",
    "        self_attention_block: MultiHeadAttentionBlock,\n",
    "        cross_attention_block: MultiHeadAttentionBlock,\n",
    "        feed_forward_block: FeedForwardBlock,\n",
    "        dropout: float,\n",
    "    ):\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 = ModuleList(\n",
    "            [ResidualConnection(dropout) for _ in range(3)]\n",
    "        )\n",
    "\n",
    "    def forward(self, x, encoder_output, src_mask, tgt_mask):\n",
    "        x = self.residual_connections[0](\n",
    "            x, lambda x: self.self_attention_block(x, x, x, tgt_mask)\n",
    "        )\n",
    "        x = self.residual_connections[1](\n",
    "            x,\n",
    "            lambda x: self.cross_attention_block(\n",
    "                x, encoder_output, encoder_output, src_mask\n",
    "            ),\n",
    "        )\n",
    "        x = self.residual_connections[2](x, self.feed_forward_block)\n",
    "        return x\n",
    "\n",
    "\n",
    "class Decoder(Module):\n",
    "    def __init__(self, layers: ModuleList):\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProjectionLayer(Module):\n",
    "    def __init__(self, d_model, vocab_size):\n",
    "        super().__init__()\n",
    "        self.proj = Linear(d_model, vocab_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, Vocab_Size)\n",
    "        return self.proj(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(Module):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        encoder: Encoder,\n",
    "        decoder: Decoder,\n",
    "        src_embed: InputEmbeddings,\n",
    "        tgt_embed: InputEmbeddings,\n",
    "        src_pos: PositionalEncoding,\n",
    "        tgt_pos: PositionalEncoding,\n",
    "        projection_layer: ProjectionLayer,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_embed = src_embed\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.src_pos = src_pos\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_transformer(\n",
    "    src_vocab_size: int,\n",
    "    tgt_vocab_size: int,\n",
    "    src_seq_len: int,\n",
    "    tgt_seq_len: int,\n",
    "    d_model: int = 512,\n",
    "    N: int = 6,\n",
    "    h: int = 8,\n",
    "    dropout: float = 0.1,\n",
    "    d_ff: int = 2048,\n",
    ") -> Transformer:\n",
    "\n",
    "    # print(f\"d_model: {type(d_model)}, src_vocab_size: {type(src_vocab_size)}\")\n",
    "    # print(f\"d_model: {d_model}, src_vocab_size: {src_vocab_size}\")\n",
    "\n",
    "    # Create the embedding layers\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 the 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, d_ff, dropout)\n",
    "        encoder_block = EncoderBlock(\n",
    "            encoder_self_attention_block, feed_forward_block, dropout\n",
    "        )\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(\n",
    "            d_model,\n",
    "            decoder_self_attention_block,\n",
    "            decoder_cross_attention_block,\n",
    "            feed_forward_block,\n",
    "            dropout,\n",
    "        )\n",
    "        decoder_blocks.append(decoder_block)\n",
    "\n",
    "    # Create the encoder and the decoder\n",
    "    encoder = Encoder(ModuleList(encoder_blocks))\n",
    "    decoder = Decoder(ModuleList(decoder_blocks))\n",
    "\n",
    "    # Create the projection layer\n",
    "    projection_layer = ProjectionLayer(d_model, tgt_vocab_size)\n",
    "\n",
    "    # Create the transformer\n",
    "    transformer = Transformer(\n",
    "        encoder, decoder, src_embed, tgt_embed, src_pos, tgt_pos, projection_layer\n",
    "    )\n",
    "\n",
    "    # Initialize the parameters\n",
    "    for parameter in transformer.parameters():\n",
    "        if parameter.dim() > 1:\n",
    "            init.xavier_uniform(parameter)\n",
    "\n",
    "    return transformer"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cloudspace",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}