Update beeper_model.py

beeper_model.py

@@ -1,22 +1,67 @@
 """
-beeper_model.py - Core model module for Beeper
-Extracted from the training code for use in inference/apps
+Rose Beeper Model V4 Fixed - Inference Components
+Extracted classes and utilities for model inference
 """
 
 import os
-import re
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from typing import Optional
-from safetensors.torch import load_file as load_safetensors
-
-# =========================================================================================
-# Model Components
-# =========================================================================================
-
+from typing import Optional, Tuple, Dict, Any
+from contextlib import nullcontext
+import re
+import inspect
+
+# ============================== Environment Setup ==============================
+torch.set_float32_matmul_precision("high")
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+# ============================== SDPA Helper ==============================
+try:
+    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:
+        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():
+    """Bias SDPA toward FlashAttention when available; no-op if unknown."""
+    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()
+
+# ============================== Model Components ==============================
 class CausalSelfAttention(nn.Module):
+    """Multi-head causal self-attention with optional FlashAttention."""
+    
     def __init__(self, dim: int, n_heads: int, attn_dropout: float = 0.0):
         super().__init__()
         assert dim % n_heads == 0
@@ -34,25 +79,34 @@ class CausalSelfAttention(nn.Module):
         k = k.view(B, T, self.nh, self.hd).transpose(1, 2)
         v = v.view(B, T, self.nh, self.hd).transpose(1, 2)
 
-        # Use scaled_dot_product_attention when available
-        y = F.scaled_dot_product_attention(
-            q, k, v,
-            is_causal=True,
-            dropout_p=self.attn_dropout if self.training else 0.0,
-        )
-        
+        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):
+    """Feed-forward MLP block with GELU activation."""
+    
     def __init__(self, dim, mlp_ratio=4.0, dropout=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):
         x = self.fc1(x)
         x = F.gelu(x, approximate="tanh")
@@ -61,23 +115,16 @@ class MLP(nn.Module):
         x = self.drop(x)
         return x
 
-
 class BeeperRoseGPT(nn.Module):
+    """Main Rose Beeper GPT model with pentachora banks."""
+    
     def __init__(self, cfg: dict):
         super().__init__()
-        V = cfg.get("vocab_size", 8192)
-        D = cfg.get("dim", 512)
-        Ctx = cfg.get("context", 512)
-        H = cfg.get("n_heads", 8)
-        L = cfg.get("n_layers", 6)
-        MR = cfg.get("mlp_ratio", 4.0)
-        RD = cfg.get("resid_dropout", 0.1)
-        AD = cfg.get("dropout", 0.0)
-
-        self.vocab_size = V
-        self.context = Ctx
-        
-        # Core transformer components
+        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"]
+
+        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)
@@ -90,25 +137,22 @@ class BeeperRoseGPT(nn.Module):
                 "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 components (for compatibility, may not be used in inference)
+        # Rose projection + anchors
         self.rose_proj = nn.Linear(D, D, bias=False)
         self.rose_anchors = nn.Parameter(torch.randn(3, D) / (D**0.5))
 
-        # Pentachora placeholders (not needed for inference but for weight compatibility)
+        # Multi-level pentachora; lazily initialized
         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.apply(self._init)
-        self.grad_checkpoint = False
+        self.grad_checkpoint = CKPT
 
     @staticmethod
     def _init(m):
@@ -119,6 +163,24 @@ 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."""
+        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)
+        self.pent_inited.fill_(1)
+
     def _block_forward(self, blk, x):
         x = x + blk["attn"](blk["norm1"](x))
         x = x + blk["mlp"](blk["norm2"](x))
@@ -128,10 +190,13 @@ class BeeperRoseGPT(nn.Module):
         B, T = idx.shape
         x = self.token_emb(idx) + self.pos_emb[:, :T, :]
         x = self.drop(x)
-        
-        for blk in self.blocks:
-            x = self._block_forward(blk, 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)
+        else:
+            for blk in self.blocks:
+                x = self._block_forward(blk, x)
         return self.norm(x)
 
     def forward(self, idx):
@@ -141,66 +206,90 @@ class BeeperRoseGPT(nn.Module):
     def hidden_states(self, idx):
         return self.backbone(idx)
 
-    def load_state_dict(self, state_dict, strict=False):
-        """Custom load that handles pentachora bank initialization gracefully"""
-        # Clean state dict keys
-        cleaned = {}
-        for k, v in state_dict.items():
+    def rose_hidden_pool(self, h: torch.Tensor, mode="mean"):
+        return h.mean(dim=1) if mode == "mean" else h[:, -1, :]
+
+# ============================== IO Utilities ==============================
+class BeeperIO:
+    """Utilities for loading and saving model checkpoints."""
+    
+    @staticmethod
+    def clean_state(sd: dict):
+        out = {}
+        for k, v in sd.items():
             if k.startswith("_orig_mod."):
                 k = k[10:]
             if k.startswith("module."):
                 k = k[7:]
-            cleaned[k] = v
+            out[k] = v
+        return out
+
+    @staticmethod
+    def load_into_model(model: nn.Module, path: str, map_location="cpu", strict: bool = False):
+        """Load weights from .pt or .safetensors file."""
+        ext = os.path.splitext(path)[1].lower()
         
-        # Initialize pentachora if present in checkpoint
-        if "penta_coarse" in cleaned:
-            self.penta_coarse = nn.Parameter(cleaned["penta_coarse"])
-        if "penta_medium" in cleaned:
-            self.penta_medium = nn.Parameter(cleaned["penta_medium"])
-        if "penta_fine" in cleaned:
-            self.penta_fine = nn.Parameter(cleaned["penta_fine"])
+        if ext == ".safetensors":
+            from safetensors.torch import load_file as load_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
         
-        return super().load_state_dict(cleaned, strict=strict)
+        sd = BeeperIO.clean_state(sd)
+        result = model.load_state_dict(sd, strict=strict)
+        return result.missing_keys, result.unexpected_keys
 
-
-# =========================================================================================
-# Generation
-# =========================================================================================
-
-def _detokenize(text: str) -> str:
-    """Clean up tokenization artifacts"""
+# ============================== Generation ==============================
+def _detok(text: str) -> str:
+    """Clean up tokenization artifacts."""
     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,  # 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: "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:
     """
-    Generate text from Beeper model with various sampling strategies.
+    Generate text from the model with various sampling strategies.
+    
+    Args:
+        model: The BeeperRoseGPT model
+        tok: Tokenizer instance
+        cfg: Configuration dictionary
+        prompt: Input prompt string
+        max_new_tokens: Maximum tokens to generate
+        temperature: Sampling temperature
+        top_k: Top-k sampling parameter
+        top_p: Top-p (nucleus) sampling parameter
+        repetition_penalty: Penalty for repeated tokens
+        presence_penalty: Penalty for token presence
+        frequency_penalty: Penalty based on token frequency
+        device: Device to run on
+        detokenize: Whether to clean up tokenization
+    
+    Returns:
+        Generated text string
     """
     # Use defaults from config if not specified
-    temperature = temperature if temperature is not None else cfg.get("temperature", 0.9)
-    top_k = top_k if top_k is not None else cfg.get("top_k", 40)
-    top_p = top_p if top_p is not None else cfg.get("top_p", 0.9)
-    repetition_penalty = repetition_penalty if repetition_penalty is not None else cfg.get("repetition_penalty", 1.1)
-    presence_penalty = presence_penalty if presence_penalty is not None else cfg.get("presence_penalty", 0.6)
-    frequency_penalty = frequency_penalty if frequency_penalty is not None else cfg.get("frequency_penalty", 0.0)
+    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
 
     device = device or next(model.parameters()).device
     model.eval()
@@ -208,19 +297,17 @@ def generate(
     # Encode prompt
     ids = tok.encode(prompt).ids
     x = torch.tensor([ids], dtype=torch.long, device=device)
+    counts = torch.zeros(cfg["vocab_size"], dtype=torch.int32, device=device)
     
-    # Track token frequencies for penalties
-    vocab_size = cfg.get("vocab_size", 8192)
-    counts = torch.zeros(vocab_size, dtype=torch.int32, device=device)
+    # Track token frequencies
     for t in ids:
-        if 0 <= t < vocab_size:
+        if 0 <= t < cfg["vocab_size"]:
             counts[t] += 1
 
     # Generate tokens
     for _ in range(max_new_tokens):
         # Get logits for next token
-        context_window = cfg.get("context", 512)
-        logits = model(x[:, -context_window:])
+        logits = model(x[:, -cfg["context"]:])
         logits = logits[:, -1, :]
 
         # Apply repetition penalty
@@ -248,24 +335,51 @@ def generate(
 
         # Top-p (nucleus) filtering
         if top_p and top_p < 1.0:
-            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-            probs = F.softmax(sorted_logits, dim=-1)
-            cumulative_probs = torch.cumsum(probs, dim=-1)
-            
-            # Find cutoff
-            cutoff_idx = (cumulative_probs > top_p).float().argmax(dim=-1)
-            mask = torch.arange(logits.size(-1), device=device).unsqueeze(0) > cutoff_idx.unsqueeze(-1)
-            sorted_logits = sorted_logits.masked_fill(mask, float("-inf"))
-            logits = torch.full_like(logits, float("-inf")).scatter(-1, sorted_indices, sorted_logits)
+            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)
 
         # Sample next token
         probs = F.softmax(logits, dim=-1)
         next_id = torch.multinomial(probs, num_samples=1)
-        
-        # Append to sequence
         x = torch.cat([x, next_id], dim=1)
         counts[next_id.item()] += 1
 
     # Decode output
-    output = tok.decode(x[0].tolist())
-    return _detokenize(output) if detokenize else output
+    out = tok.decode(x[0].tolist())
+    return _detok(out) if detokenize else out
+
+# ============================== Default Configuration ==============================
+def get_default_config():
+    """Return the default configuration for the Rose Beeper 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 parameters
+        "temperature": 0.9,
+        "top_k": 40,
+        "top_p": 0.9,
+        "repetition_penalty": 1.10,
+        "presence_penalty": 0.6,
+        "frequency_penalty": 0.0,
+        
+        # Capoera settings
+        "capoera": {
+            "enable": True,
+            "topic_bins": 512,
+            "mood_bins": 7,
+        }
+    }