model.py

import math

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


# ---------------- Masks ----------------
def make_padding_mask(token_ids: torch.Tensor, pad_id: int = 0) -> torch.Tensor:
    """
    token_ids: (B, T)
    returns: (B,1,1,T) where True means BLOCK (masked out)
    """
    return (token_ids == pad_id).unsqueeze(1).unsqueeze(2)

def make_causal_mask(T: int, device: torch.device) -> torch.Tensor:
    """
    returns: (1,1,T,T) where True above diagonal means BLOCK future positions
    """
    return torch.triu(torch.ones(T, T, dtype=torch.bool, device=device), diagonal=1).unsqueeze(0).unsqueeze(0)

# ---------------- Core Attention ----------------
def scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, training=True):
    """
    q,k,v: (B,H,T,Dh)
    attn_mask: broadcastable to (B,H,T,T) where True=BLOCK
    """
    Dh = q.size(-1)
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(Dh)  # (B,H,T,T)
    if attn_mask is not None:
        scores = scores.masked_fill(attn_mask, -1e9)
    attn = F.softmax(scores, dim=-1)
    if dropout_p > 0:
        attn = F.dropout(attn, p=dropout_p, training=training)
    out = attn @ v
    return out, attn


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model=128, n_heads=4, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.dropout = dropout

        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, attn_mask=None):
        """
        x: (B,T,d_model)
        attn_mask: broadcastable to (B,H,T,T) where True=BLOCK
        returns:
          out: (B,T,d_model)
          attn: (B,H,T,T)
        """
        B, T, _ = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)

        def split(t):
            return t.view(B, T, self.n_heads, self.d_head).transpose(1, 2)  # (B,H,T,Dh)

        q = split(q)
        k = split(k)
        v = split(v)

        out, attn = scaled_dot_product_attention(q, k, v, attn_mask, self.dropout, self.training)
        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        out = self.proj(out)
        return out, attn
    

class LearnablePositionalEmbedding(nn.Module):
    def __init__(self, max_len: int, d_model: int):
        super().__init__()
        self.pos = nn.Embedding(max_len, d_model)

    def forward(self, x):
        """
        x: (B,T,d_model)
        """
        T = x.size(1)
        positions = torch.arange(T, device=x.device).unsqueeze(0)  # (1,T)
        return x + self.pos(positions)


class FeedForward(nn.Module):
    def __init__(self, d_model=128, d_ff=512, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )

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


class GPTDecoderBlock(nn.Module):
    def __init__(self, d_model=128, n_heads=4, d_ff=512, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadSelfAttention(d_model, n_heads, dropout)
        self.drop1 = nn.Dropout(dropout)

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

    def forward(self, x, attn_mask=None):
        a_out, attn = self.attn(self.ln1(x), attn_mask=attn_mask)
        x = x + self.drop1(a_out)

        f_out = self.ffn(self.ln2(x))
        x = x + self.drop2(f_out)
        return x, attn
    
class LM_GPT(nn.Module):
    def __init__(self, vocab_size: int, d_model=128, n_heads=4, n_layers=2, max_len=256, d_ff=512, dropout=0.1, pad_id=0):
        super().__init__()
        self.pad_id = pad_id
        self.emb = nn.Embedding(vocab_size, d_model, padding_idx=pad_id)
        self.pos = LearnablePositionalEmbedding(max_len=max_len, d_model=d_model)
        self.blocks = nn.ModuleList([GPTDecoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, token_ids: torch.Tensor):
        """
        token_ids: (B,T)
        returns:
          logits: (B,T,V)
          attn_last: (B,H,T,T)
        """
        x = self.pos(self.emb(token_ids))
        B, T = token_ids.shape
        causal = make_causal_mask(T, token_ids.device)
        pad_mask = make_padding_mask(token_ids, self.pad_id)
        attn_mask = causal | pad_mask  # broadcastable to (B,H,T,T)

        attn_last = None
        for blk in self.blocks:
            x, attn_last = blk(x, attn_mask=attn_mask)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        return logits, attn_last