Update beeper_model.py

beeper_model.py

@@ -1,33 +1,38 @@
-"""
-Rose Beeper Model - Inference Components
-Extracted classes and utilities for model inference
-"""
+# 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 os
 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
-from typing import Optional, Tuple, Dict, Any
-from contextlib import nullcontext
-import inspect
-import re
-from tokenizers import Tokenizer
-from safetensors.torch import load_file as load_safetensors
-
 
-# ============================================================================
-# SDPA (Scaled Dot Product Attention) Configuration
-# ============================================================================
+# --- 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 context helper
+# ---- 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
@@ -39,23 +44,23 @@ except Exception:
 
 
 def sdpa_ctx_prefer_flash():
-    """Bias SDPA toward FlashAttention when available; no-op if unknown."""
+    """
+    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:
-        # Modern API (PyTorch 2.3+): backends=[...]
         if "backends" in params and _SDPBackend is not None:
             return _sdpa_kernel(backends=[
                 _SDPBackend.FLASH_ATTENTION,
                 _SDPBackend.EFFICIENT_ATTENTION,
                 _SDPBackend.MATH
             ])
-        # Modern API (alt): backend=...
         if "backend" in params and _SDPBackend is not None:
             return _sdpa_kernel(backend=_SDPBackend.FLASH_ATTENTION)
-        # Legacy boolean flags (old CUDA backend)
         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:
@@ -65,27 +70,27 @@ def sdpa_ctx_prefer_flash():
     return nullcontext()
 
 
-# ============================================================================
-# Model Components
-# ============================================================================
-
+# --------------------------------- Core blocks ------------------------------------------------------
 class CausalSelfAttention(nn.Module):
-    """Multi-head causal self-attention with optional FlashAttention."""
-    
+    """
+    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
-        self.nh = n_heads
-        self.hd = dim // n_heads
+        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 = attn_dropout
+        self.attn_dropout = float(attn_dropout)
 
-    def forward(self, x):
+    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)
+        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)
 
@@ -109,16 +114,15 @@ class CausalSelfAttention(nn.Module):
 
 
 class MLP(nn.Module):
-    """Feed-forward network with GELU activation."""
-    
-    def __init__(self, dim, mlp_ratio=4.0, dropout=0.1):
+    """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):
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.fc1(x)
         x = F.gelu(x, approximate="tanh")
         x = self.drop(x)
@@ -127,47 +131,62 @@ class MLP(nn.Module):
         return x
 
 
+# --------------------------------- Beeper Model -----------------------------------------------------
 class BeeperRoseGPT(nn.Module):
-    """Rose Beeper GPT model with pentachora banks for multi-level control."""
-    
+    """
+    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, CKPT = cfg["resid_dropout"], cfg["dropout"], cfg["grad_checkpoint"]
+        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.vocab_size, self.context = V, Ctx
         self.token_emb = nn.Embedding(V, D)
-        self.pos_emb = nn.Parameter(torch.zeros(1, Ctx, D))
-        self.drop = nn.Dropout(RD)
+        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),
+                "attn":  CausalSelfAttention(D, H, attn_dropout=AD),
                 "norm2": nn.LayerNorm(D),
-                "mlp": MLP(D, mlp_ratio=MR, dropout=RD),
-            }) for _ in range(L)
+                "mlp":   MLP(D, mlp_ratio=MR, dropout=RD),
+            })
+            for _ in range(L)
         ])
-        self.norm = nn.LayerNorm(D)
+
+        self.norm    = nn.LayerNorm(D)
         self.lm_head = nn.Linear(D, V, bias=False)
+
+        # Weight tying
         self.lm_head.weight = self.token_emb.weight
 
-        # Optional Rose projection + anchors
-        self.rose_proj = nn.Linear(D, D, bias=False)
-        self.rose_anchors = nn.Parameter(torch.randn(3, D) / (D**0.5))
+        # 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))
 
-        # Multi-level pentachora; lazily initialized
+        # Pentachora banks (created lazily to match state dict)
         self.register_buffer("pent_inited", torch.tensor(0, dtype=torch.uint8), persistent=False)
-        self.penta_coarse = None
-        self.penta_medium = None
-        self.penta_fine = None
+        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)
-        self.grad_checkpoint = CKPT
+        self.apply(self._init_weights)
 
     @staticmethod
-    def _init(m):
+    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:
@@ -175,92 +194,95 @@ class BeeperRoseGPT(nn.Module):
         elif isinstance(m, nn.Embedding):
             nn.init.normal_(m.weight, mean=0.0, std=0.02)
 
-    def ensure_pentachora(self, coarse_C: int, medium_C: int, fine_C: int, dim: int, device):
-        """Initialize three pentachora banks."""
+    # ---- 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):
-            pts = []
-            for _ in range(int(C)):
-                A = torch.randn(5, dim, device=device)
-                A = F.normalize(A - A.mean(dim=0, keepdim=True), dim=-1)
-                pts.append(A)
-            return nn.Parameter(torch.stack(pts, dim=0))
-
-        self.penta_coarse = bank(coarse_C)
-        self.penta_medium = bank(medium_C)
-        self.penta_fine = bank(fine_C)
+        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)
 
-    def _block_forward(self, blk, x):
+    # ---- 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):
+    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)
+                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):
+    def forward(self, idx: torch.Tensor) -> torch.Tensor:
         h = self.backbone(idx)
         return self.lm_head(h)
 
-    def hidden_states(self, idx):
+    # ---- 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="mean"):
+    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, :]
 
 
-# ============================================================================
-# Model I/O Utilities
-# ============================================================================
+# --------------------------------- 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.
 
-class BeeperIO:
-    """Utilities for saving and loading model weights."""
-    
-    @staticmethod
-    def clean_state(sd: dict):
-        """Clean state dict keys from various wrappings."""
-        out = {}
-        for k, v in sd.items():
-            if k.startswith("_orig_mod."):
-                k = k[10:]
-            if k.startswith("module."):
-                k = k[7:]
-            out[k] = v
-        return out
+    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
 
-    @staticmethod
-    def load_into_model(model: nn.Module, path: str, map_location="cpu", strict: bool = False):
-        """Load weights from file into model."""
-        ext = os.path.splitext(path)[1].lower()
-        if ext == ".safetensors":
-            sd = load_safetensors(path, device="cpu")
-        else:
-            raw = torch.load(path, map_location="cpu")
-            sd = raw["model"] if isinstance(raw, dict) and "model" in raw else raw
-        sd = BeeperIO.clean_state(sd)
-        result = model.load_state_dict(sd, strict=strict)
-        return result.missing_keys, result.unexpected_keys
+    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)
 
-# ============================================================================
-# Text Generation
-# ============================================================================
 
+# --------------------------------- Generation -------------------------------------------------------
 def _detok(text: str) -> str:
-    """Clean up tokenized text spacing."""
     text = re.sub(r"\s+([,.;:!?%])", r"\1", text)
     text = re.sub(r"\s+([\)\]\}])", r"\1", text)
     text = re.sub(r"([\(\[\{])\s+", r"\1", text)
@@ -268,92 +290,67 @@ def _detok(text: str) -> str:
 
 
 @torch.no_grad()
-def generate(model: BeeperRoseGPT, 
-             tok: Tokenizer, 
-             cfg: dict, 
-             prompt: str,
-             max_new_tokens: int = 120, 
-             temperature: float = None, 
-             top_k: int = None, 
-             top_p: float = None,
-             repetition_penalty: float = None, 
-             presence_penalty: float = None, 
-             frequency_penalty: float = None,
-             device: Optional[torch.device] = None, 
-             detokenize: bool = True) -> str:
+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:
     """
-    Generate text from a prompt using the model.
-    
-    Args:
-        model: The BeeperRoseGPT model
-        tok: Tokenizer instance
-        cfg: Configuration dictionary
-        prompt: Input text prompt
-        max_new_tokens: Maximum number of tokens to generate
-        temperature: Sampling temperature (higher = more random)
-        top_k: Top-k sampling parameter
-        top_p: Top-p (nucleus) sampling parameter
-        repetition_penalty: Penalty for repeated tokens
-        presence_penalty: Penalty for tokens that have appeared
-        frequency_penalty: Penalty based on token frequency
-        device: Device to run on
-        detokenize: Whether to clean up tokenization artifacts
-    
-    Returns:
-        Generated text string
+    Penalized nucleus sampling (same knobs as training script).
     """
-    
-    # Use defaults from config if not specified
-    temperature = cfg["temperature"] if temperature is None else temperature
-    top_k = cfg["top_k"] if top_k is None else top_k
-    top_p = cfg["top_p"] if top_p is None else top_p
-    repetition_penalty = cfg["repetition_penalty"] if repetition_penalty is None else repetition_penalty
-    presence_penalty = cfg["presence_penalty"] if presence_penalty is None else presence_penalty
-    frequency_penalty = cfg["frequency_penalty"] if frequency_penalty is None else frequency_penalty
+    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()
-    
-    # Tokenize prompt
+
     ids = tok.encode(prompt).ids
     x = torch.tensor([ids], dtype=torch.long, device=device)
-    
-    # Track token counts for penalties
-    counts = torch.zeros(cfg["vocab_size"], dtype=torch.int32, device=device)
+    V = int(cfg["vocab_size"])
+    counts = torch.zeros(V, dtype=torch.int32, device=device)
     for t in ids:
-        if 0 <= t < cfg["vocab_size"]:
+        if 0 <= t < V:
             counts[t] += 1
 
-    # Generate tokens
-    for _ in range(max_new_tokens):
-        # Get logits for next token
+    for _ in range(int(max_new_tokens)):
         logits = model(x[:, -cfg["context"]:])
         logits = logits[:, -1, :]
 
-        # Apply repetition penalty
+        # 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
                 logits[:, mask][~pos] *= repetition_penalty
 
-        # Apply presence and frequency penalties
+        # 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)
 
-        # Apply temperature
         logits = logits / max(1e-8, temperature)
 
-        # Apply top-k sampling
         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)
 
-        # Apply top-p (nucleus) sampling
         if top_p and top_p < 1.0:
             sl, si = torch.sort(logits, descending=True)
             ps = F.softmax(sl, dim=-1)
@@ -363,47 +360,12 @@ def generate(model: BeeperRoseGPT,
             sl = sl.masked_fill(mask, float("-inf"))
             logits = torch.full_like(logits, float("-inf")).scatter(-1, si, sl)
 
-        # Sample next token
         probs = F.softmax(logits, dim=-1)
         next_id = torch.multinomial(probs, num_samples=1)
         x = torch.cat([x, next_id], dim=1)
-        counts[next_id.item()] += 1
+        nid = next_id.item()
+        if 0 <= nid < V:
+            counts[nid] += 1
 
-    # Decode output
     out = tok.decode(x[0].tolist())
     return _detok(out) if detokenize else out
-
-
-# ============================================================================
-# Default Configuration
-# ============================================================================
-
-def get_default_config():
-    """Get the default configuration for the model."""
-    return {
-        "name": "Rose-Beeper",
-        "context": 512,
-        "vocab_size": 8192,
-        "dim": 512,
-        "n_layers": 6,
-        "n_heads": 8,
-        "mlp_ratio": 4.0,
-        "dropout": 0.0,
-        "resid_dropout": 0.1,
-        "grad_checkpoint": False,
-        
-        # Generation defaults
-        "temperature": 0.9,
-        "top_k": 40,
-        "top_p": 0.9,
-        "repetition_penalty": 1.10,
-        "presence_penalty": 0.6,
-        "frequency_penalty": 0.0,
-        
-        # Capoera configuration
-        "capoera": {
-            "enable": True,
-            "topic_bins": 512,
-            "mood_bins": 7,
-        }
-    }