model.py

import torch.nn as nn
from einops import einsum, rearrange
import torch 
import math




def softmax(x , dim):
    x_max = x.max(dim=dim, keepdim=True).values
    exp = torch.exp(x-x_max)

    return exp/torch.sum(exp, dim=dim, keepdim=True)


def scaled_dot_product_attention(q, k, v, mask=None):
    QK = einsum(q, k, "... seq_len d_k, ... seq d_k -> ... seq_len seq")
    d_k = q.shape[-1]
    root_d_k = 1/math.sqrt(d_k)

    s = QK * root_d_k
    if mask is not None:
        s = s.masked_fill(mask == False, float('-inf'))
    s_max = softmax(s, -1)

    return einsum(s_max, v, "... q k, ... k d_v -> ... q d_v")



class Linear(nn.Module):
    def __init__(self, in_features: int, out_features: int, device = None, dtype = None):
        super().__init__()

        self.weight = nn.Parameter(
            torch.empty(out_features, in_features, device = device, dtype = dtype)
            )
        nn.init.trunc_normal_(
            self.weight,
              mean = 0.0,
                std = math.sqrt(2/(in_features + out_features)), 
                a = -3.0 * math.sqrt(2/(in_features + out_features)),
                  b = 3.0 * math.sqrt(2/(in_features + out_features))
            )
    def forward(self, A):
        return einsum(A, self.weight, "... d_in, d_out d_in -> ... d_out")



class Embedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, device = None, dtype = None):
        super().__init__()
        self.weight = nn.Parameter(
            torch.empty(num_embeddings, embedding_dim, device=device, dtype=dtype)
        )
        nn.init.trunc_normal_(self.weight, mean = 0.0, std = 1.0, a = -3.0, b = 3.0)

    def forward(self, token_ids):
        return self.weight[token_ids]



class RMSNorm(nn.Module):
    def __init__(self, d_model, eps: float = 1e-5 , device = None, dtype = None):
        super().__init__()
        self.weight = nn.Parameter(
            torch.ones(d_model) 
        )
        self.eps = eps
        # No trunc normal since we're not drawing from a distribution.

    def forward(self, x):
        in_dtype = x.dtype
        x = x.to(torch.float32) # convert to float32 to avoid overflow
        sq = x*x
        m = sq.mean(dim=-1, keepdim= True)
        result = x * torch.rsqrt(m + self.eps) * self.weight
        return result.to(in_dtype)



class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = Linear(d_model, d_ff)
        self.w2 = Linear(d_ff, d_model)
        self.w3 = Linear(d_model, d_ff)

    def forward(self, x):
        w1x = self.w1(x)
        w3x = self.w3(x)
        siluw1x = w1x * torch.sigmoid(w1x)

        return self.w2(siluw1x * w3x)


class SiLU(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x * torch.sigmoid(x)



# RoPE
class RotaryPositionalEmbedding(nn.Module):
    def __init__(self, theta: float, d_k: int, max_seq_len: int, device = None):
        super().__init__()
        self.theta = theta
        self.d_k = d_k
        self.max_seq_len = max_seq_len
        value_vector = torch.arange(0, d_k, 2).float()
        frequency_vector = theta ** (-(value_vector/d_k))
        positions_vector = torch.arange(max_seq_len).float()
        angles = torch.outer(positions_vector, frequency_vector)

        cos_table = angles.cos()
        sin_table = angles.sin()

        self.register_buffer("cos_table", cos_table, persistent=False)
        self.register_buffer("sin_table", sin_table, persistent=False)

    def forward(self, x, token_positions):
        pairs = rearrange(x, '... seq (half two) -> ... seq half two', two = 2) # Get a 2d vector of pairs. 
        a = pairs[..., 0]
        b = pairs[..., 1]

        cos = self.cos_table[token_positions]
        sin = self.sin_table[token_positions]

        a_rot = a * cos - b * sin
        b_rot = a * sin + b * cos

        stacked = torch.stack((a_rot, b_rot), dim = -1)

        return rearrange(stacked, '... seq half two -> ... seq (half two)')





class Multihead_attention(nn.Module):
    def __init__(self, d_model, num_heads, max_seq_len = None, theta = None, do_rope: bool = None):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.W_K = Linear(d_model, d_model)
        self.W_Q = Linear(d_model, d_model)
        self.W_V = Linear(d_model, d_model)
        self.W_O = Linear(d_model, d_model)
        self.do_rope = do_rope
        if self.do_rope == True:
            self.rope = RotaryPositionalEmbedding(theta, self.d_k, max_seq_len)
        

    def forward(self, x):
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)

        Q = rearrange(Q, "batch seq (h d_k) -> batch h seq d_k", h=self.num_heads)
        K = rearrange(K, "batch seq (h d_k) -> batch h seq d_k", h=self.num_heads)
        V = rearrange(V, "batch seq (h d_k) -> batch h seq d_k", h=self.num_heads)

        # Apply RoPE.
        seq = x.shape[1]
        token_positions = torch.arange(seq, device= x.device)
        if self.do_rope == True:
            Q = self.rope(Q, token_positions)
            K = self.rope(K, token_positions)

        # Apply causal masking
        mask = torch.tril(torch.ones(seq, seq, dtype=torch.bool, device = x.device))


        out = scaled_dot_product_attention(Q, K, V, mask)


        out = rearrange(out, "batch h seq d_k -> batch seq (h d_k)")

        return self.W_O(out)



class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, max_seq_len, theta):
        super().__init__()
        self.attn = Multihead_attention(d_model, num_heads, max_seq_len, theta, do_rope = True)
        self.ffn = SwiGLU(d_model, d_ff)
        self.norm1 = RMSNorm(d_model)
        self.norm2 = RMSNorm(d_model)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x



class TransformerLM(nn.Module):
    def __init__(self, vocab_size, context_length, num_layers, d_model, num_heads, d_ff, theta):
        super().__init__()
        self.embedding = Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList(
            [
                TransformerBlock(d_model, num_heads, d_ff, context_length, theta)
                for _ in range(num_layers)
            ]
        )
        self.final_norm = RMSNorm(d_model)
        self.lm_head = Linear(d_model, vocab_size)

    def forward(self, token_ids):
        x = self.embedding(token_ids)
        for layer in self.layers:
            x = layer(x)
        x = self.final_norm(x)
        return self.lm_head(x)