decoder.py

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


# ============================================================
# Causal Mask
# ============================================================

def causal_mask(seq_len, device):
    """
    Prevent attending to future tokens.

    shape: (1, 1, seq_len, seq_len)
    """
    mask = torch.tril(torch.ones(seq_len, seq_len, device=device))
    mask = mask.unsqueeze(0).unsqueeze(0)
    return mask


# ============================================================
# Multi-Head Self Attention
# ============================================================

class MultiHeadSelfAttention(nn.Module):

    def __init__(self, d_model, n_heads):
        super().__init__()

        assert d_model % n_heads == 0

        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads

        # combined projection (faster, used in real models)
        self.qkv = nn.Linear(d_model, 3 * d_model)

        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x, mask):

        B, T, C = x.shape

        qkv = self.qkv(x)  # (B, T, 3C)

        qkv = qkv.view(B, T, 3, self.n_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)

        Q, K, V = qkv[0], qkv[1], qkv[2]

        # attention scores
        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.head_dim)

        scores = scores.masked_fill(mask[:, :, :T, :T] == 0, float('-inf'))

        attn = torch.softmax(scores, dim=-1)

        out = attn @ V

        out = out.transpose(1, 2).contiguous().view(B, T, C)

        out = self.out_proj(out)

        return out


# ============================================================
# Feedforward Network
# ============================================================

class FeedForward(nn.Module):

    def __init__(self, d_model, d_ff):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model)
        )

    def forward(self, x):
        return self.net(x)


# ============================================================
# Decoder Block
# ============================================================

class DecoderBlock(nn.Module):

    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()

        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadSelfAttention(d_model, n_heads)

        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FeedForward(d_model, d_ff)

    def forward(self, x, mask):

        # attention
        x = x + self.attn(self.ln1(x), mask)

        # feedforward
        x = x + self.ffn(self.ln2(x))

        return x


# ============================================================
# Full Decoder-Only Transformer
# ============================================================

class DecoderOnlyTransformer(nn.Module):

    def __init__(
        self,
        vocab_size,
        d_model=768,
        n_heads=12,
        n_layers=12,
        d_ff=3072,
        max_seq_len=1024
    ):
        super().__init__()

        self.token_embedding = nn.Embedding(vocab_size, d_model)

        self.position_embedding = nn.Embedding(max_seq_len, d_model)

        self.blocks = nn.ModuleList([
            DecoderBlock(d_model, n_heads, d_ff)
            for _ in range(n_layers)
        ])

        self.ln_final = nn.LayerNorm(d_model)

        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

        self.max_seq_len = max_seq_len

    def forward(self, tokens):

        B, T = tokens.shape

        device = tokens.device

        pos = torch.arange(0, T, device=device)
        pos = pos.unsqueeze(0)

        x = self.token_embedding(tokens) + self.position_embedding(pos)

        mask = causal_mask(self.max_seq_len, device)

        for block in self.blocks:
            x = block(x, mask)

        x = self.ln_final(x)

        logits = self.lm_head(x)

        return logits


# ============================================================
# Test run
# ============================================================

if __name__ == "__main__":

    model = DecoderOnlyTransformer(
        vocab_size=50000,
        d_model=512,
        n_heads=8,
        n_layers=6,
        d_ff=2048,
        max_seq_len=256
    )

    x = torch.randint(0, 50000, (2, 128))

    logits = model(x)

    print(logits.shape)

    # expected:
    # (batch, seq_len, vocab_size)


### code written by chat --- just a reference implementation