decoder_block.py

"""
decoder_block.py

Builds a complete Transformer Decoder Block by combining:

1. Multi-Head Self Attention
2. Residual Connection
3. Layer Normalization
4. Feed Forward Network
5. Residual Connection
6. Layer Normalization

This is the fundamental building block used in decoder-only
Large Language Models such as GPT, Llama, Mistral, Gemma,
Qwen, Phi and DeepSeek.
"""

import torch
import torch.nn as nn

from attention import MultiHeadSelfAttention
from feed_forward import FeedForward
from config import DROPOUT


class DecoderBlock(nn.Module):
    """
    A single Transformer Decoder Block.

    Architecture:

        Input
          │
          ▼
    Multi-Head Self Attention
          │
    Add (Residual)
          │
    LayerNorm
          │
    Feed Forward Network
          │
    Add (Residual)
          │
    LayerNorm
          │
        Output
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        ff_hidden_dim: int,
    ):
        super().__init__()

        # Multi-Head Self Attention
        self.attention = MultiHeadSelfAttention(
            embed_dim,
            num_heads,
        )

        # Feed Forward Network
        self.feed_forward = FeedForward(
            embed_dim,
            ff_hidden_dim,
        )

        # Layer Normalization
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)

        # Dropout
        self.dropout = nn.Dropout(DROPOUT)

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        # -------------------------------
        # Multi-Head Self Attention
        # -------------------------------
        attention_output = self.attention(x)

        # Residual Connection + LayerNorm
        x = self.norm1(
            x + self.dropout(attention_output)
        )

        # -------------------------------
        # Feed Forward Network
        # -------------------------------
        ff_output = self.feed_forward(x)

        # Residual Connection + LayerNorm
        x = self.norm2(
            x + self.dropout(ff_output)
        )

        return x