Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,311 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math
+from dataclasses import dataclass
+from contextlib import nullcontext
+from typing import Literal 
+
+
+class CausalSelfAttention(nn.Module):
+    # A causal self-attention layer that supports both flash attention and standard attention.
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0 # Ensures the embedding dimension can be evenly split across attention heads.
+        
+        # This linear layer projects input x into query (q), key (k), and value (v) vectors — 
+        # all at once (so the output is 3× the size).
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        
+        # After attention is done, this layer projects the output back to the original embedding size.
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        
+        # Dropout applied to the attention weights (probabilities).
+        self.attn_dropout = nn.Dropout(config.dropout)
+        
+        # Dropout applied after the final projection.
+        self.resid_dropout = nn.Dropout(config.dropout)
+        
+        # Store values for easy access later.
+        self.n_head = config.n_head 
+        self.n_embd = config.n_embd
+        
+        # Checks whether the efficient Flash Attention API is available in torch.nn.functional.
+        self.flash = hasattr(F, "scaled_dot_product_attention")
+        
+        # If Flash Attention is not available, we create a lower triangular mask to ensure causality.
+        # This mask prevents the model from attending to future tokens in the sequence.
+        if not self.flash:
+            # register_buffer ensures this tensor is saved with the model but not updated by gradients.
+            self.register_buffer(
+                "bias",
+                torch.tril(torch.ones(config.block_size, config.block_size)).view(
+                    1, 1, config.block_size, config.block_size
+                ),
+            )
+    def forward(self, x):
+        B, T, C = x.size()
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
+        if self.flash:
+            y = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=None,
+                dropout_p=self.attn_dropout.p if self.training else 0.0,
+                is_causal=True,
+            )
+        else:
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf"))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+# --- User's Original LayerNorm ---
+class LayerNorm(nn.Module):
+    def __init__(self, ndim, bias):
+        """
+        Initializes the LayerNorm module.
+        Args:
+            ndim (int): is the number of features in the last dimension (e.g., embedding size).
+            bias (bool): Whether to include a bias term in the normalization.
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, x):
+        return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5)
+# --- End User's Original LayerNorm ---
+
+# --- User's Original MLP ---
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu = nn.GELU()
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))
+# --- End User's Original MLP ---
+
+# --- User's Original Block ---
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln1 = LayerNorm(config.n_embd, config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln2 = LayerNorm(config.n_embd, config.bias)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+# --- End User's Original Block ---
+
+
+# --- User's Original GPTConfig ---
+@dataclass
+class GPTConfig:
+    block_size: int
+    vocab_size: int
+    n_layer: int
+    n_head: int
+    n_embd: int
+    dropout: float = 0.0
+    bias: bool = True
+# --- End User's Original GPTConfig ---
+
+# --- User's Original TrainingConfig ---
+@dataclass
+class TrainingConfig:
+    learning_rate: float = 1e-4  # more stable training, earlier 1e-4
+    max_iters: int = 20000  # increase from 25000
+    warmup_steps: int = 1000  # smoother initial train, earlier 100
+    min_lr: float = 5e-4  # lower rate, earlier 5e-4
+    eval_iters: int = 500  # increased from 100
+    batch_size: int = 32  # changed from 16, better gradient estimate
+    block_size: int = 128  # changed from 64, capture longer range dependencies
+    gradient_accumulation_steps: int = 32  # reduced from 50
+    device: Literal["cuda", "cpu"] = "cuda" if torch.cuda.is_available() else "cpu"
+    device_type: Literal["cuda", "cpu"] = (
+        "cuda" if "cuda" in device else "cpu"
+    )  # for later use in torch.autocast
+    dtype: Literal["bfloat16", "float16"] = (
+        "bfloat16"
+        if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+        else "float16"
+    )
+    ptdtype: torch.dtype = {
+        "float32": torch.float32,
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+    }[dtype]
+    ctx: nullcontext[None] | torch.autocast = (
+        nullcontext()
+        if device_type == "cpu"
+        else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+    )
+# --- End User's Original TrainingConfig ---
+
+
+class GPT(nn.Module):
+    """
+    The main GPT model, now with an optional QA head for Question Answering tasks.
+    The QA head will predict start and end token indices of the answer span.
+    """
+    def __init__(self, config, is_qa_model=False):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+        self.is_qa_model = is_qa_model
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = LayerNorm(config.n_embd, bias=config.bias),
+        ))
+        
+        # Language modeling head (for pre-training)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        
+        # QA head (for fine-tuning)
+        # This will predict start and end logits for the answer span
+        if self.is_qa_model:
+            self.qa_head = nn.Linear(config.n_embd, 2, bias=False) # 2 outputs: start_logit, end_logit
+        else:
+            self.qa_head = None # No QA head if not a QA model
+
+        # tie weights
+        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
+
+        # init all weights
+        self.apply(self._init_weights)
+        # apply special scaled init to the residual projections, per GPT-2 paper
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/((2 * config.n_layer)**0.5))
+
+        # report number of parameters
+        # n_params calculation will differ slightly if QA head is present
+        n_params = sum(p.numel() for p in self.parameters())
+        # For non-embedding count it excludes token embeddings and positional embeddings.
+        non_embedding_params = n_params - self.transformer.wpe.weight.numel()
+        print(f"Number of parameters: {non_embedding_params/1e6:.2f}M (excluding positional embeddings)")
+
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, input_ids, targets=None, attention_mask=None, token_type_ids=None):
+        device = input_ids.device
+        b, t = input_ids.size()
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(input_ids) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if self.is_qa_model and self.qa_head is not None:
+            # For QA, we typically use the pooled output or sequence output directly
+            # For extractive QA, we need logits for each token for start/end prediction
+            # The output 'x' is (batch_size, sequence_length, n_embd)
+            logits = self.qa_head(x) # (batch_size, sequence_length, 2)
+            start_logits, end_logits = logits.split(1, dim=-1)
+            start_logits = start_logits.squeeze(-1).contiguous() # (batch_size, sequence_length)
+            end_logits = end_logits.squeeze(-1).contiguous()     # (batch_size, sequence_length)
+
+            if targets is not None:
+                # targets for QA are start_positions and end_positions
+                start_positions, end_positions = targets[:, 0], targets[:, 1]
+                
+                # Apply attention mask to logits for valid tokens
+                if attention_mask is not None:
+                    # Tokens that are part of the context (token_type_ids == 1) should be considered for answers
+                    # and also non-padding tokens (attention_mask == 1)
+                    valid_tokens_mask = (attention_mask == 1) & (token_type_ids == 1)
+                    
+                    start_logits = start_logits.masked_fill(~valid_tokens_mask, float('-inf'))
+                    end_logits = end_logits.masked_fill(~valid_tokens_mask, float('-inf'))
+
+                loss_fct = nn.CrossEntropyLoss(ignore_index=-100) # Use -100 as ignore_index for consistency
+                start_loss = loss_fct(start_logits, start_positions)
+                end_loss = loss_fct(end_logits, end_positions)
+                total_loss = (start_loss + end_loss) / 2
+                return start_logits, end_logits, total_loss
+            
+            return start_logits, end_logits, None # For inference
+        else: # Standard language model for pre-training or text generation
+            if targets is not None:
+                # if we are given some targets (e.g. for training), calculate the loss
+                logits = self.lm_head(x)
+                loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100) # Use -100
+            else:
+                # inference-time mini-optimization: only forward the lm_head on the very last position
+                logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
+                loss = None
+
+            return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        """
+        Generate tokens given a conditioning sequence.
+        idx: Tensor of shape (B, T)
+        """
+        if self.is_qa_model:
+            print("Warning: generate method is not intended for QA models directly.")
+            print("Please use the QA forward pass for inference and post-processing.")
+            return idx # Or raise an error
+            
+        for _ in range(max_new_tokens):
+            idx_cond = (
+                idx
+                if idx.size(1) <= self.config.block_size
+                else idx[:, -self.config.block_size :]
+            )
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("Inf")
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx
+
+# The 'config' object for pre-training is also kept here, if it's used by other scripts for its definition
+config = GPTConfig(
+    vocab_size=50257,  # use the tokenizer's vocab size
+    block_size=1024,   # or whatever context size you're training with
+    n_layer=8,
+    n_head=8,
+    n_embd=512,
+    dropout=0.1,
+    bias=True,
+)