Upload architecture.py with huggingface_hub

architecture.py

@@ -0,0 +1,125 @@
+# 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.
+# ============================================================================
+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):
+    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
+
+        # 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.
+        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)
+
+        # 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)
+        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).
+        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:
+            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().
+        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):
+    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.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(d_ff, d_model)   # project: 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):
+    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.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, is_causal=True):
+        # Sub-layer 1: normalize -> attention -> dropout -> ADD back to input.
+        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