model.py

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,
)