Update beeper_model.py

beeper_model.py

@@ -1,5 +1,5 @@
 """
-Rose Beeper Model V4 Fixed - Inference Components
+Rose Beeper Model - Inference Components
 Extracted classes and utilities for model inference
 """
 
@@ -10,15 +10,17 @@ import torch.nn as nn
 import torch.nn.functional as F
 from typing import Optional, Tuple, Dict, Any
 from contextlib import nullcontext
-import re
 import inspect
+import re
+from tokenizers import Tokenizer
+from safetensors.torch import load_file as load_safetensors
+
 
-# ============================== Environment Setup ==============================
-torch.set_float32_matmul_precision("high")
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
+# ============================================================================
+# SDPA (Scaled Dot Product Attention) Configuration
+# ============================================================================
 
-# ============================== SDPA Helper ==============================
+# Version-safe SDPA context helper
 try:
     from torch.nn.attention import sdpa_kernel as _sdpa_kernel_modern
     from torch.nn.attention import SDPBackend as _SDPBackend
@@ -35,6 +37,7 @@ except Exception:
         _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:
@@ -42,14 +45,17 @@ def sdpa_ctx_prefer_flash():
 
     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:
@@ -58,7 +64,11 @@ def sdpa_ctx_prefer_flash():
         pass
     return nullcontext()
 
-# ============================== Model Components ==============================
+
+# ============================================================================
+# Model Components
+# ============================================================================
+
 class CausalSelfAttention(nn.Module):
     """Multi-head causal self-attention with optional FlashAttention."""
     
@@ -97,8 +107,9 @@ class CausalSelfAttention(nn.Module):
         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."""
+    """Feed-forward network with GELU activation."""
     
     def __init__(self, dim, mlp_ratio=4.0, dropout=0.1):
         super().__init__()
@@ -115,8 +126,9 @@ class MLP(nn.Module):
         x = self.drop(x)
         return x
 
+
 class BeeperRoseGPT(nn.Module):
-    """Main Rose Beeper GPT model with pentachora banks."""
+    """Rose Beeper GPT model with pentachora banks for multi-level control."""
     
     def __init__(self, cfg: dict):
         super().__init__()
@@ -141,7 +153,7 @@ class BeeperRoseGPT(nn.Module):
         self.lm_head = nn.Linear(D, V, bias=False)
         self.lm_head.weight = self.token_emb.weight
 
-        # Rose projection + anchors
+        # 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))
 
@@ -209,12 +221,17 @@ class BeeperRoseGPT(nn.Module):
     def rose_hidden_pool(self, h: torch.Tensor, mode="mean"):
         return h.mean(dim=1) if mode == "mean" else h[:, -1, :]
 
-# ============================== IO Utilities ==============================
+
+# ============================================================================
+# Model I/O Utilities
+# ============================================================================
+
 class BeeperIO:
-    """Utilities for loading and saving model checkpoints."""
+    """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."):
@@ -226,31 +243,33 @@ class BeeperIO:
 
     @staticmethod
     def load_into_model(model: nn.Module, path: str, map_location="cpu", strict: bool = False):
-        """Load weights from .pt or .safetensors file."""
+        """Load weights from file into model."""
         ext = os.path.splitext(path)[1].lower()
-        
         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
-        
         sd = BeeperIO.clean_state(sd)
         result = model.load_state_dict(sd, strict=strict)
         return result.missing_keys, result.unexpected_keys
 
-# ============================== Generation ==============================
+
+# ============================================================================
+# Text Generation
+# ============================================================================
+
 def _detok(text: str) -> str:
-    """Clean up tokenization artifacts."""
+    """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)
     return text
 
+
 @torch.no_grad()
 def generate(model: BeeperRoseGPT, 
-             tok: "Tokenizer", 
+             tok: Tokenizer, 
              cfg: dict, 
              prompt: str,
              max_new_tokens: int = 120, 
@@ -263,26 +282,27 @@ def generate(model: BeeperRoseGPT,
              device: Optional[torch.device] = None, 
              detokenize: bool = True) -> str:
     """
-    Generate text from the model with various sampling strategies.
+    Generate text from a prompt using the model.
     
     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
+        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 token presence
+        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
+        detokenize: Whether to clean up tokenization artifacts
     
     Returns:
         Generated text string
     """
+    
     # 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
@@ -294,12 +314,12 @@ def generate(model: BeeperRoseGPT,
     device = device or next(model.parameters()).device
     model.eval()
     
-    # Encode prompt
+    # Tokenize 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
+    # Track token counts for penalties
+    counts = torch.zeros(cfg["vocab_size"], dtype=torch.int32, device=device)
     for t in ids:
         if 0 <= t < cfg["vocab_size"]:
             counts[t] += 1
@@ -323,17 +343,17 @@ def generate(model: BeeperRoseGPT,
             pen = counts.float() * (frequency_penalty or 0.0) + (counts > 0).float() * (presence_penalty or 0.0)
             logits = logits - pen.unsqueeze(0)
 
-        # Temperature scaling
+        # Apply temperature
         logits = logits / max(1e-8, temperature)
 
-        # Top-k filtering
+        # 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)
 
-        # Top-p (nucleus) filtering
+        # 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)
@@ -353,9 +373,13 @@ def generate(model: BeeperRoseGPT,
     out = tok.decode(x[0].tolist())
     return _detok(out) if detokenize else out
 
-# ============================== Default Configuration ==============================
+
+# ============================================================================
+# Default Configuration
+# ============================================================================
+
 def get_default_config():
-    """Return the default configuration for the Rose Beeper model."""
+    """Get the default configuration for the model."""
     return {
         "name": "Rose-Beeper",
         "context": 512,
@@ -368,7 +392,7 @@ def get_default_config():
         "resid_dropout": 0.1,
         "grad_checkpoint": False,
         
-        # Generation parameters
+        # Generation defaults
         "temperature": 0.9,
         "top_k": 40,
         "top_p": 0.9,
@@ -376,7 +400,7 @@ def get_default_config():
         "presence_penalty": 0.6,
         "frequency_penalty": 0.0,
         
-        # Capoera settings
+        # Capoera configuration
         "capoera": {
             "enable": True,
             "topic_bins": 512,