model.py

import torch
import torch.nn as nn
import torch.nn.functional as F

from typing import Dict


# embedding layer
class EmbeddingLayer(nn.Module):
    def __init__(self, vocab_size, max_context, embed_dim):
        super().__init__()
        self.semantic_embedding = nn.Embedding(vocab_size, embed_dim)
        self.positional_embedding = nn.Embedding(max_context, embed_dim)

    def forward(self, x):
        # indices are of shape (B, T)
        x1 = self.semantic_embedding(x)

        # positions a list of ints from 0 to block_size-1
        _, max_context = x.shape
        positions = torch.arange(max_context, device=x.device)
        x2 = self.positional_embedding(positions)

        # output are of shape (B, T, embed_dim)
        x = x1 + x2

        return x


# self-attention layer
class SelfAttentionLayer(nn.Module):
    def __init__(self, embed_dim, head_size, dropout):
        super().__init__()
        self.query = nn.Linear(embed_dim, head_size, bias=False)
        self.key = nn.Linear(embed_dim, head_size, bias=False)
        self.value = nn.Linear(embed_dim, head_size, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # compute q, k, and v
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)

        # construct attention scores matrix
        scores = q @ k.transpose(1, 2)

        # get shape and scale scores
        _, T, A = q.shape
        scores = scores / (A ** 0.5)

        # causal masking (decoder only)
        mask = torch.ones(T, T, device=x.device)
        mask = torch.tril(mask)
        mask = mask == 0
        scores = scores.masked_fill(mask, float('-inf'))

        # generate output of shape (B, T, head_size)
        scores = F.softmax(scores, dim=-1)
        scores = self.dropout(scores)
        x = scores @ v

        return x


# multi-head attention
class MHA(nn.Module):
    def __init__(self, embed_dim, attention_dim, num_heads, dropout):
        super().__init__()
        head_size = attention_dim // num_heads
        self.mha = nn.ModuleList()
        for _ in range(num_heads):
            self.mha.append(SelfAttentionLayer(embed_dim, head_size, dropout))
        self.projection = nn.Linear(head_size * num_heads, embed_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        results = []
        for head in self.mha:
            results.append(head(x))
        x = torch.cat(results, dim=-1)
        x = self.projection(x)
        x = self.dropout(x)

        return x


# feed-forward network
class FFN(nn.Module):
    def __init__(self, embed_dim, ffn_dim, dropout):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, ffn_dim),
            nn.ReLU()
        )
        self.projection = nn.Linear(ffn_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        x = self.mlp(x)
        x = self.projection(x)
        x = self.dropout(x)

        return x


# transformer block
class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, attention_dim, num_heads, ffn_dim, dropout):
        super().__init__()
        self.mha = MHA(embed_dim, attention_dim, num_heads, dropout)
        self.ffn = FFN(embed_dim, ffn_dim, dropout)
        self.ln1 = nn.LayerNorm(embed_dim)
        self.ln2 = nn.LayerNorm(embed_dim)

    def forward(self, x):
        # residual connections
        x = x + self.mha(self.ln1(x))
        x = x + self.ffn(self.ln2(x))

        return x
        
        
# GPT model
class GPT(nn.Module):
    def __init__(self, vocab_size, max_context, embed_dim, attention_dim, num_heads, ffn_dim, dropout, num_blocks):
        super().__init__()
        self.embed = EmbeddingLayer(vocab_size, max_context, embed_dim)
        self.blocks = nn.Sequential()
        for _ in range(num_blocks):
            self.blocks.append(TransformerBlock(embed_dim, attention_dim, num_heads, ffn_dim, dropout))
        self.ln = nn.LayerNorm(embed_dim)
        self.projection = nn.Linear(embed_dim, vocab_size)

        self.max_context = max_context

    def forward(self, x):
        x = self.embed(x)
        x = self.blocks(x)
        x = self.ln(x)
        x = self.projection(x)

        # return logits (not probabilities)
        return x