model.py

"""
model.py

Builds a minimal Decoder-Only Transformer Language Model by stacking
multiple Transformer Decoder Blocks.

Pipeline:

Token IDs
    ↓
Token Embedding + Positional Encoding
    ↓
N × Decoder Blocks
    ↓
Final Layer Normalization
    ↓
Linear Projection
    ↓
Vocabulary Logits
"""

import torch
import torch.nn as nn

from embeddings import TokenEmbedding
from decoder_block import DecoderBlock

from config import (
    VOCAB_SIZE,
    EMBED_DIM,
    NUM_HEADS,
    FF_HIDDEN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
)


class MiniDecoderLM(nn.Module):
    """
    A minimal Decoder-Only Transformer Language Model.

    This project is designed for learning the internal architecture
    of modern decoder-only Large Language Models (LLMs) such as
    GPT, Llama, Mistral, Gemma, Qwen, Phi, and DeepSeek.
    """

    def __init__(self):
        super().__init__()

        # Token Embedding + Positional Encoding
        self.embedding = TokenEmbedding(
            vocab_size=VOCAB_SIZE,
            embed_dim=EMBED_DIM,
            max_len=MAX_SEQ_LEN,
        )

        # Stack of Decoder Blocks
        self.decoder_blocks = nn.ModuleList(
            [
                DecoderBlock(
                    embed_dim=EMBED_DIM,
                    num_heads=NUM_HEADS,
                    ff_hidden_dim=FF_HIDDEN_DIM,
                )
                for _ in range(NUM_LAYERS)
            ]
        )

        # Final Layer Normalization
        self.final_norm = nn.LayerNorm(EMBED_DIM)

        # Output Projection
        self.output_projection = nn.Linear(
            EMBED_DIM,
            VOCAB_SIZE,
        )

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

        # Convert token IDs into embeddings
        x = self.embedding(token_ids)

        # Pass through all Decoder Blocks
        for decoder in self.decoder_blocks:
            x = decoder(x)

        # Final normalization
        x = self.final_norm(x)

        # Project to vocabulary
        logits = self.output_projection(x)

        return logits