Upload architecture.py with huggingface_hub

architecture.py

@@ -1,125 +1,78 @@
-# architecture.py
-# ============================================================================
-# The Transformer building blocks. A Transformer block = self-attention (tokens
-# share information) + a feed-forward network (each token is processed on its
-# own), each wrapped in a residual connection and LayerNorm. These three classes
-# are stacked by LogSentryLM in train_and_generate.py.
-# ============================================================================
+# file: architecture.py
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import math
 
 
-# ============================================================================
-# MULTI-HEAD SELF-ATTENTION
-# Lets each token look at other tokens and pull in the information it needs.
-# Each token forms a Query ("what am I looking for?"), a Key ("what do I
-# offer?"), and a Value ("what I'll pass on"). A token attends to others by
-# matching its Query against their Keys, then takes a weighted sum of Values.
-# ============================================================================
 class MultiHeadAttention(nn.Module):
+    """Scaled dot-product self-attention with a fused QKV projection."""
+
     def __init__(self, d_model, n_heads):
         super().__init__()
-        # d_model must split evenly across the heads, since each head works on a
-        # d_model/n_heads slice of the vector.
         assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
         self.d_model = d_model
         self.n_heads = n_heads
-        self.d_k = d_model // n_heads   # dimension PER head
+        self.d_k = d_model // n_heads
 
-        # One linear layer produces Q, K and V at once (hence 3 * d_model), then
-        # we split the result. Doing it in a single matmul is faster than three.
+        # One matmul for Q, K and V is cheaper than three separate projections
         self.qkv_projection = nn.Linear(d_model, 3 * d_model, bias=False)
-        # Mixes the heads' outputs back together into one d_model vector.
         self.out_projection = nn.Linear(d_model, d_model, bias=False)
 
     def forward(self, x, is_causal=True):
-        # x: (batch_size, seq_len, d_model)
         batch_size, seq_len, d_model = x.shape
-
-        # Project to Q,K,V combined, then split along the last dim into three.
         qkv = self.qkv_projection(x)
-        q, k, v = qkv.chunk(3, dim=-1)   # each: (batch, seq_len, d_model)
+        q, k, v = qkv.chunk(3, dim=-1)
 
-        # Reshape each into separate heads and move the head axis forward so the
-        # math runs independently per head:
-        # (batch, seq_len, d_model) -> (batch, n_heads, seq_len, d_k)
+        # (B, T, d_model) -> (B, heads, T, d_k) so each head runs independently
         q = q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
         k = k.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
         v = v.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
 
-        # Attention scores = how much each token should attend to each other
-        # token = Q . K^T. Dividing by sqrt(d_k) keeps the numbers from getting
-        # too large (which would make softmax too "peaky" and hurt gradients).
+        # Scale by sqrt(d_k) — keeps the scores from blowing up and the softmax stable
         scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
 
-        # CAUSAL MASK: for next-token prediction, a token may only attend to
-        # itself and EARLIER tokens — never the future (that would be cheating).
-        # tril keeps the lower triangle; the rest is set to -inf so that after
-        # softmax those positions get weight ~0.
         if is_causal:
+            # Block attention to future tokens — required for next-token prediction,
+            # otherwise the model would just read the answer.
             mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device)).view(1, 1, seq_len, seq_len)
             scores = scores.masked_fill(mask == 0, float('-inf'))
 
-        # Turn scores into attention weights that sum to 1 across the keys.
         attention_weights = F.softmax(scores, dim=-1)
-
-        # Weighted sum of Values = the information each token gathers.
         context = torch.matmul(attention_weights, v)
 
-        # Merge the heads back together: (batch, n_heads, seq_len, d_k) ->
-        # (batch, seq_len, d_model). .contiguous() makes memory layout valid for
-        # the following view().
+        # Merge heads back: (B, heads, T, d_k) -> (B, T, d_model)
         context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
-
-        # Final linear mix across the combined heads.
         return self.out_projection(context)
 
 
-# ============================================================================
-# FEED-FORWARD NETWORK
-# After attention shares information ACROSS tokens, this processes EACH token
-# independently. It expands to a wider hidden size, applies a non-linearity,
-# then projects back — giving the model room to compute richer features.
-# ============================================================================
 class FeedForwardNetwork(nn.Module):
+    """Per-token MLP: expand to d_ff, non-linearity, project back."""
+
     def __init__(self, d_model, d_ff, dropout=0.1):
         super().__init__()
-        # d_ff is the wider inner dimension. Convention is d_ff = 4 * d_model
-        # (set by the caller), giving the network extra capacity to "think".
-        self.linear1 = nn.Linear(d_model, d_ff)   # expand:  d_model -> d_ff
-        # Dropout randomly zeros some activations during TRAINING only, which
-        # reduces overfitting. dropout=0.1 means 10% are dropped.
+        self.linear1 = nn.Linear(d_model, d_ff)
         self.dropout = nn.Dropout(dropout)
-        self.linear2 = nn.Linear(d_ff, d_model)   # project: d_ff -> d_model
+        self.linear2 = nn.Linear(d_ff, d_model)
 
     def forward(self, x):
-        # GELU is a smooth activation function (a standard choice in Transformers).
         return self.linear2(self.dropout(F.gelu(self.linear1(x))))
 
 
-# ============================================================================
-# TRANSFORMER BLOCK
-# Combines attention + feed-forward with two key tricks:
-#   * Residual connections (x + sublayer(x)): help gradients flow and let each
-#     layer make small refinements instead of rebuilding from scratch.
-#   * Pre-norm (LayerNorm BEFORE each sublayer): the modern GPT-2 style; trains
-#     more stably than the original "post-norm" design.
-# ============================================================================
 class TransformerBlock(nn.Module):
+    """One decoder block: attention + feed-forward, pre-norm with residuals."""
+
     def __init__(self, d_model, n_heads, dropout=0.1):
         super().__init__()
         self.attention = MultiHeadAttention(d_model, n_heads)
-        # d_ff = 4 * d_model is the standard feed-forward expansion ratio.
-        self.feed_forward = FeedForwardNetwork(d_model, d_ff=4 * d_model, dropout=dropout)
-        self.ln1 = nn.LayerNorm(d_model)   # normalizes input to the attention
-        self.ln2 = nn.LayerNorm(d_model)   # normalizes input to the feed-forward
+        self.feed_forward = FeedForwardNetwork(d_model, d_ff=4 * d_model, dropout=dropout)  # 4x is the usual ratio
+        self.ln1 = nn.LayerNorm(d_model)
+        self.ln2 = nn.LayerNorm(d_model)
         self.dropout = nn.Dropout(dropout)
 
     def forward(self, x, is_causal=True):
-        # Sub-layer 1: normalize -> attention -> dropout -> ADD back to input.
+        # Pre-norm (LayerNorm before each sub-layer) trains more stably than post-norm.
+        # The `x + ...` residuals let gradients flow and each layer refine rather than replace.
         x = x + self.dropout(self.attention(self.ln1(x), is_causal=is_causal))
-        # Sub-layer 2: normalize -> feed-forward -> dropout -> ADD back to input.
         x = x + self.dropout(self.feed_forward(self.ln2(x)))
         return x