beeper_model.py

# beeper.py
# --------------------------------------------------------------------------------------------------
# Beeper — Rose-based tiny GPT (inference module)
# - Decoder-only GPT with SDPA (FlashAttention path on Ampere+)
# - Model exactly mirrors the training-time architecture you provided (dim=512, L=6, H=8)
# - Safe state-dict loader that auto-sizes pentachora banks before strict load
# - Generation API with repetition/presence/frequency penalties (same defaults as training)
# --------------------------------------------------------------------------------------------------
from __future__ import annotations

import math
import re
import inspect
from contextlib import nullcontext
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

# --- Prefer high-throughput matmul where possible (Ampere/Hopper) ---
torch.set_float32_matmul_precision("high")
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# ---- Version-safe SDPA (FlashAttention) selection -------------------------------------------------
try:
    # PyTorch 2.3+ modern API
    from torch.nn.attention import sdpa_kernel as _sdpa_kernel_modern
    from torch.nn.attention import SDPBackend as _SDPBackend
    _SDPA_SIG = inspect.signature(_sdpa_kernel_modern)
    _sdpa_kernel = _sdpa_kernel_modern
except Exception:
    try:
        # Legacy API
        from torch.backends.cuda import sdp_kernel as _sdpa_kernel_legacy
        _SDPA_SIG = inspect.signature(_sdpa_kernel_legacy)
        _SDPBackend = None
        _sdpa_kernel = _sdpa_kernel_legacy
    except Exception:
        _SDPA_SIG = None
        _SDPBackend = None
        _sdpa_kernel = None


def sdpa_ctx_prefer_flash():
    """
    Best-effort context to bias SDPA toward FlashAttention on supported GPUs.
    Falls back to no-op if not available.
    """
    if _sdpa_kernel is None or _SDPA_SIG is None:
        return nullcontext()

    params = {p.name for p in _SDPA_SIG.parameters.values()}
    try:
        if "backends" in params and _SDPBackend is not None:
            return _sdpa_kernel(backends=[
                _SDPBackend.FLASH_ATTENTION,
                _SDPBackend.EFFICIENT_ATTENTION,
                _SDPBackend.MATH
            ])
        if "backend" in params and _SDPBackend is not None:
            return _sdpa_kernel(backend=_SDPBackend.FLASH_ATTENTION)
        if {"enable_flash", "enable_math", "enable_mem_efficient"} <= params:
            return _sdpa_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=True)
        if {"use_flash", "use_math", "use_mem_efficient"} <= params:
            return _sdpa_kernel(use_flash=True, use_math=False, use_mem_efficient=True)
    except Exception:
        pass
    return nullcontext()


# --------------------------------- Core blocks ------------------------------------------------------
class CausalSelfAttention(nn.Module):
    """
    Multi-head causal self-attention layer using PyTorch SDPA.
    - On CUDA, uses scaled_dot_product_attention with is_causal=True and dropout during training.
    - On CPU, falls back to manual masked attention.
    """
    def __init__(self, dim: int, n_heads: int, attn_dropout: float = 0.0):
        super().__init__()
        assert dim % n_heads == 0, "dim must be divisible by n_heads"
        self.nh = int(n_heads)
        self.hd = dim // self.nh
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.proj = nn.Linear(dim, dim, bias=False)
        self.attn_dropout = float(attn_dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
        q = q.view(B, T, self.nh, self.hd).transpose(1, 2)  # [B,H,T,D]
        k = k.view(B, T, self.nh, self.hd).transpose(1, 2)
        v = v.view(B, T, self.nh, self.hd).transpose(1, 2)

        if x.is_cuda:
            with sdpa_ctx_prefer_flash():
                y = F.scaled_dot_product_attention(
                    q, k, v,
                    is_causal=True,
                    dropout_p=self.attn_dropout if self.training else 0.0,
                )
        else:
            scale = 1.0 / math.sqrt(self.hd)
            att = (q @ k.transpose(-2, -1)) * scale
            mask = torch.full((1, 1, T, T), float("-inf"), device=x.device)
            mask = torch.triu(mask, diagonal=1)
            att = (att + mask).softmax(dim=-1)
            y = att @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(y)


class MLP(nn.Module):
    """GELU MLP with dropout, sized by mlp_ratio."""
    def __init__(self, dim: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
        super().__init__()
        hidden = int(dim * mlp_ratio)
        self.fc1 = nn.Linear(dim, hidden)
        self.fc2 = nn.Linear(hidden, dim)
        self.drop = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc1(x)
        x = F.gelu(x, approximate="tanh")
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


# --------------------------------- Beeper Model -----------------------------------------------------
class BeeperRoseGPT(nn.Module):
    """
    Decoder-only GPT used by Beeper during training and inference.

    Config keys used:
      - vocab_size, dim, context, n_heads, n_layers, mlp_ratio
      - resid_dropout, dropout, grad_checkpoint
    Notes:
      - Shares token embedding with LM head (tied weights).
      - Includes Rose projection/anchors and pentachora banks; unused for plain generation,
        but kept for full compatibility with trained checkpoints.
    """
    def __init__(self, cfg: dict):
        super().__init__()
        V, D, Ctx = cfg["vocab_size"], cfg["dim"], cfg["context"]
        H, L, MR = cfg["n_heads"], cfg["n_layers"], cfg["mlp_ratio"]
        RD, AD = cfg.get("resid_dropout", 0.1), cfg.get("dropout", 0.0)
        self.grad_checkpoint = bool(cfg.get("grad_checkpoint", False))

        self.vocab_size, self.context = int(V), int(Ctx)

        self.token_emb = nn.Embedding(V, D)
        self.pos_emb   = nn.Parameter(torch.zeros(1, Ctx, D))
        self.drop      = nn.Dropout(RD)

        self.blocks = nn.ModuleList([
            nn.ModuleDict({
                "norm1": nn.LayerNorm(D),
                "attn":  CausalSelfAttention(D, H, attn_dropout=AD),
                "norm2": nn.LayerNorm(D),
                "mlp":   MLP(D, mlp_ratio=MR, dropout=RD),
            })
            for _ in range(L)
        ])

        self.norm    = nn.LayerNorm(D)
        self.lm_head = nn.Linear(D, V, bias=False)

        # Weight tying
        self.lm_head.weight = self.token_emb.weight

        # Rose projection + anchors (present in checkpoints)
        self.rose_proj    = nn.Linear(D, D, bias=False)
        self.rose_anchors = nn.Parameter(torch.randn(3, D) / (D ** 0.5))

        # Pentachora banks (created lazily to match state dict)
        self.register_buffer("pent_inited", torch.tensor(0, dtype=torch.uint8), persistent=False)
        self.penta_coarse: Optional[nn.Parameter] = None  # [C,5,D]
        self.penta_medium: Optional[nn.Parameter] = None  # [T,5,D]
        self.penta_fine:   Optional[nn.Parameter] = None  # [M,5,D]

        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(m: nn.Module):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    # ---- Pentachora creation (must match sizes in checkpoint before strict load) -------------------
    def ensure_pentachora(self, coarse_C: int, medium_C: int, fine_C: int, dim: int, device: torch.device):
        """
        Initialize pentachora banks if not already present.
        Shapes must match checkpoint entries for strict loading.
        """
        if self.pent_inited.item() == 1:
            return

        def bank(C: int) -> nn.Parameter:
            if C <= 0:
                # Keep a zero-sized parameter to satisfy strict loading (rare).
                return nn.Parameter(torch.zeros((0, 5, dim), device=device))
            pts = torch.randn(C, 5, dim, device=device)
            pts = F.normalize(pts - pts.mean(dim=1, keepdim=True), dim=-1)
            return nn.Parameter(pts)

        self.penta_coarse = bank(int(coarse_C))
        self.penta_medium = bank(int(medium_C))
        self.penta_fine   = bank(int(fine_C))
        self.pent_inited.fill_(1)

    # ---- Backbone / forward -----------------------------------------------------------------------
    def _block_forward(self, blk: nn.ModuleDict, x: torch.Tensor) -> torch.Tensor:
        x = x + blk["attn"](blk["norm1"](x))
        x = x + blk["mlp"](blk["norm2"](x))
        return x

    def backbone(self, idx: torch.Tensor) -> torch.Tensor:
        B, T = idx.shape
        x = self.token_emb(idx) + self.pos_emb[:, :T, :]
        x = self.drop(x)
        if self.grad_checkpoint and self.training:
            from torch.utils.checkpoint import checkpoint
            for blk in self.blocks:
                x = checkpoint(lambda _x: self._block_forward(blk, _x), x)  # type: ignore[arg-type]
        else:
            for blk in self.blocks:
                x = self._block_forward(blk, x)
        return self.norm(x)

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        h = self.backbone(idx)
        return self.lm_head(h)

    # ---- Utilities ---------------------------------------------------------------------------------
    def hidden_states(self, idx: torch.Tensor) -> torch.Tensor:
        """Return final hidden states (pre-LM head)."""
        return self.backbone(idx)

    def rose_hidden_pool(self, h: torch.Tensor, mode: str = "mean") -> torch.Tensor:
        """Pool hidden states for Rose-related terms (unused in plain generation)."""
        return h.mean(dim=1) if mode == "mean" else h[:, -1, :]


# --------------------------------- Loader helpers ---------------------------------------------------
def prepare_model_for_state_dict(
    model: BeeperRoseGPT,
    state_dict: "dict[str, torch.Tensor]",
    device: Optional[torch.device] = None,
) -> None:
    """
    Ensure model has pentachora parameters sized to match the incoming state_dict,
    so we can load with strict=True.

    If the checkpoint has no pentachora (older versions), we do nothing.
    """
    device = device or next(model.parameters()).device
    need = all(k in state_dict for k in ("penta_coarse", "penta_medium", "penta_fine"))
    if not need:
        return

    pc, pt, pm = state_dict["penta_coarse"], state_dict["penta_medium"], state_dict["penta_fine"]
    # Expect [C,5,D]
    def dims_ok(t: torch.Tensor) -> bool:
        return t.ndim == 3 and t.size(1) == 5 and t.size(2) == model.token_emb.embedding_dim

    if not (dims_ok(pc) and dims_ok(pt) and dims_ok(pm)):
        # Shapes inconsistent; fall back to non-strict load later.
        return

    coarse_C = pc.size(0)
    topic_C  = pt.size(0)
    mood_C   = pm.size(0)
    model.ensure_pentachora(coarse_C, topic_C, mood_C, dim=pc.size(2), device=device)


# --------------------------------- Generation -------------------------------------------------------
def _detok(text: str) -> str:
    text = re.sub(r"\s+([,.;:!?%])", r"\1", text)
    text = re.sub(r"\s+([\)\]\}])", r"\1", text)
    text = re.sub(r"([\(\[\{])\s+", r"\1", text)
    return text


@torch.no_grad()
def generate(
    model: BeeperRoseGPT,
    tok,                      # Hugging Face Tokenizers `Tokenizer`
    cfg: dict,
    prompt: str,
    max_new_tokens: int = 120,
    temperature: Optional[float] = None,
    top_k: Optional[int] = None,
    top_p: Optional[float] = None,
    repetition_penalty: Optional[float] = None,
    presence_penalty: Optional[float] = None,
    frequency_penalty: Optional[float] = None,
    device: Optional[torch.device] = None,
    detokenize: bool = True,
) -> str:
    """
    Penalized nucleus sampling (same knobs as training script).
    """
    temperature = cfg.get("temperature", 0.9) if temperature is None else float(temperature)
    top_k = cfg.get("top_k", 40) if top_k is None else int(top_k)
    top_p = cfg.get("top_p", 0.9) if top_p is None else float(top_p)
    repetition_penalty = cfg.get("repetition_penalty", 1.10) if repetition_penalty is None else float(repetition_penalty)
    presence_penalty = cfg.get("presence_penalty", 0.6) if presence_penalty is None else float(presence_penalty)
    frequency_penalty = cfg.get("frequency_penalty", 0.0) if frequency_penalty is None else float(frequency_penalty)

    device = device or next(model.parameters()).device
    model.eval()

    ids = tok.encode(prompt).ids
    x = torch.tensor([ids], dtype=torch.long, device=device)
    V = int(cfg["vocab_size"])
    counts = torch.zeros(V, dtype=torch.int32, device=device)
    for t in ids:
        if 0 <= t < V:
            counts[t] += 1

    for _ in range(int(max_new_tokens)):
        logits = model(x[:, -cfg["context"]:])
        logits = logits[:, -1, :]

        # Repetition penalty (CTRL-like)
        if repetition_penalty and repetition_penalty != 1.0:
            mask = counts > 0
            if mask.any():
                pos = logits[:, mask] > 0
                logits[:, mask][pos]  /= repetition_penalty
                logits[:, mask][~pos] *= repetition_penalty

        # Presence/frequency penalties (OpenAI-like)
        if presence_penalty or frequency_penalty:
            pen = counts.float() * (frequency_penalty or 0.0) + (counts > 0).float() * (presence_penalty or 0.0)
            logits = logits - pen.unsqueeze(0)

        logits = logits / max(1e-8, temperature)

        if top_k and top_k > 0:
            k = min(top_k, logits.size(-1))
            v, ix = torch.topk(logits, k, dim=-1)
            filt = torch.full_like(logits, float("-inf"))
            logits = filt.scatter_(-1, ix, v)

        if top_p and top_p < 1.0:
            sl, si = torch.sort(logits, descending=True)
            ps = F.softmax(sl, dim=-1)
            cdf = torch.cumsum(ps, dim=-1)
            cutoff = (cdf > top_p).float().argmax(dim=-1)
            mask = torch.arange(logits.size(-1), device=device).unsqueeze(0) > cutoff.unsqueeze(-1)
            sl = sl.masked_fill(mask, float("-inf"))
            logits = torch.full_like(logits, float("-inf")).scatter(-1, si, sl)

        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        x = torch.cat([x, next_id], dim=1)
        nid = next_id.item()
        if 0 <= nid < V:
            counts[nid] += 1

    out = tok.decode(x[0].tolist())
    return _detok(out) if detokenize else out