Update beeper_model.py

beeper_model.py

@@ -1,10 +1,10 @@
 # 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)
+# Beeper Full Penta Controller — Rose-based tiny GPT (inference module with runtime pentachora influence)
+# - Decoder-only GPT with SDPA (FlashAttention path on Ampere/Hopper)
+# - Runtime "vertex pull" uses config["runtime_pentachora"] to bias hidden states toward
+#   pentachora vertices (coarse/topic/mood) exactly like training-time behavior, but non-destructive
+#   and fully toggleable.
 # --------------------------------------------------------------------------------------------------
 from __future__ import annotations
 
@@ -12,7 +12,7 @@ import math
 import re
 import inspect
 from contextlib import nullcontext
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Dict, Any
 
 import torch
 import torch.nn as nn
@@ -44,13 +44,9 @@ except Exception:
 
 
 def sdpa_ctx_prefer_flash():
-    """
-    Best-effort context to bias SDPA toward FlashAttention on supported GPUs.
-    Falls back to no-op if not available.
-    """
+    """Bias SDPA toward FlashAttention where possible; otherwise no-op."""
     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:
@@ -72,11 +68,7 @@ def sdpa_ctx_prefer_flash():
 
 # --------------------------------- 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.
-    """
+    """Multi-head causal self-attention using PyTorch SDPA."""
     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"
@@ -139,10 +131,18 @@ class BeeperRoseGPT(nn.Module):
     Config keys used:
       - vocab_size, dim, context, n_heads, n_layers, mlp_ratio
       - resid_dropout, dropout, grad_checkpoint
+      - runtime_pentachora: {
+            "enable": bool,
+            "pool": "mean" | "last",
+            "temp": float,               # similarity temperature (default: 0.10)
+            "coarse_alpha": float,       # hidden blend strength for coarse bank
+            "topic_alpha":  float,       # hidden blend strength for topic bank
+            "mood_alpha":   float        # hidden blend strength for mood bank
+        }
     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.
+      - Includes Rose anchors and pentachora banks; at runtime we can apply a *non-destructive*
+        vertex pull to hidden states before the LM head using the above config.
     """
     def __init__(self, cfg: dict):
         super().__init__()
@@ -150,6 +150,7 @@ class BeeperRoseGPT(nn.Module):
         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.runtime_cfg: Dict[str, Any] = dict(cfg.get("runtime_pentachora", {}) or {})
 
         self.vocab_size, self.context = int(V), int(Ctx)
 
@@ -169,9 +170,7 @@ class BeeperRoseGPT(nn.Module):
 
         self.norm    = nn.LayerNorm(D)
         self.lm_head = nn.Linear(D, V, bias=False)
-
-        # Weight tying
-        self.lm_head.weight = self.token_emb.weight
+        self.lm_head.weight = self.token_emb.weight  # weight tying
 
         # Rose projection + anchors (present in checkpoints)
         self.rose_proj    = nn.Linear(D, D, bias=False)
@@ -196,16 +195,12 @@ class BeeperRoseGPT(nn.Module):
 
     # ---- 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.
-        """
+        """Initialize pentachora banks if not already present."""
         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)
@@ -216,6 +211,94 @@ class BeeperRoseGPT(nn.Module):
         self.penta_fine   = bank(int(fine_C))
         self.pent_inited.fill_(1)
 
+    # ---- Runtime configuration helpers -------------------------------------------------------------
+    def set_runtime_pentachora(self, cfg: Dict[str, Any]) -> None:
+        """Update runtime pentachora behavior (enable/alphas/temp/pool)."""
+        self.runtime_cfg.update(cfg or {})
+
+    def _pool_hidden(self, h: torch.Tensor, mode: str) -> torch.Tensor:
+        return h.mean(dim=1) if mode == "mean" else h[:, -1, :]
+
+    @staticmethod
+    def _weighted_nearest_vertex_target(
+        pooled: torch.Tensor,          # [B,D]
+        bank: torch.Tensor,            # [C,5,D]
+        temp: float
+    ) -> torch.Tensor:
+        """
+        For each class (simplex) pick its nearest vertex to the pooled latent,
+        then compute a softmax over classes of -min_dists/temp and take the
+        weighted average of those nearest vertices => [B,D] target.
+        """
+        B, D = pooled.shape
+        C = bank.size(0)
+        if C == 0:
+            return pooled
+
+        # distances to each vertex
+        diffs = pooled[:, None, None, :] - bank[None, :, :, :]    # [B,C,5,D]
+        dists = torch.norm(diffs, dim=-1)                         # [B,C,5]
+
+        min_dists, min_idx = dists.min(dim=2)                     # [B,C], [B,C]
+        sims = -min_dists / max(1e-8, float(temp))                # [B,C]
+        weights = F.softmax(sims, dim=-1)                         # [B,C]
+
+        # gather nearest vertex vectors: [B,C,D]
+        bank_exp = bank.unsqueeze(0).expand(B, -1, -1, -1)        # [B,C,5,D]
+        gather_idx = min_idx.unsqueeze(-1).unsqueeze(-1).expand(B, C, 1, D)
+        nearest = torch.gather(bank_exp, 2, gather_idx).squeeze(2)  # [B,C,D]
+
+        target = (weights.unsqueeze(-1) * nearest).sum(dim=1)     # [B,D]
+        return target
+
+    def _apply_runtime_vertex_pull(
+        self,
+        h: torch.Tensor,                        # [B,T,D]
+        runtime_cfg: Dict[str, Any]
+    ) -> torch.Tensor:
+        """
+        Apply non-destructive vertex pull to hidden states using banks selected by runtime_cfg.
+        We compute a pooled latent, a per-bank target vector, form a delta, and blend it back into h.
+        """
+        if not runtime_cfg or not runtime_cfg.get("enable", False):
+            return h
+
+        pool_mode = str(runtime_cfg.get("pool", "mean"))
+        temp = float(runtime_cfg.get("temp", 0.10))
+
+        # Strengths per bank
+        alpha_coarse = float(runtime_cfg.get("coarse_alpha", 0.0))
+        alpha_topic  = float(runtime_cfg.get("topic_alpha",  0.0))
+        alpha_mood   = float(runtime_cfg.get("mood_alpha",   0.0))
+
+        if (alpha_coarse <= 0 and alpha_topic <= 0 and alpha_mood <= 0):
+            return h
+
+        pooled = self._pool_hidden(h, pool_mode)                  # [B,D]
+
+        total_delta = None
+        if alpha_coarse > 0 and getattr(self, "penta_coarse", None) is not None:
+            tgt = self._weighted_nearest_vertex_target(pooled, self.penta_coarse, temp)
+            delta = tgt - pooled
+            total_delta = (alpha_coarse * delta) if total_delta is None else total_delta + alpha_coarse * delta
+
+        if alpha_topic > 0 and getattr(self, "penta_medium", None) is not None:
+            tgt = self._weighted_nearest_vertex_target(pooled, self.penta_medium, temp)
+            delta = tgt - pooled
+            total_delta = delta * alpha_topic if total_delta is None else total_delta + alpha_topic * delta
+
+        if alpha_mood > 0 and getattr(self, "penta_fine", None) is not None:
+            tgt = self._weighted_nearest_vertex_target(pooled, self.penta_fine, temp)
+            delta = tgt - pooled
+            total_delta = delta * alpha_mood if total_delta is None else total_delta + alpha_mood * delta
+
+        if total_delta is None:
+            return h
+
+        # Broadcast same delta to all time steps (global conditioning shift)
+        h = h + total_delta.unsqueeze(1)       # [B,T,D]
+        return h
+
     # ---- Backbone / forward -----------------------------------------------------------------------
     def _block_forward(self, blk: nn.ModuleDict, x: torch.Tensor) -> torch.Tensor:
         x = x + blk["attn"](blk["norm1"](x))
@@ -235,8 +318,14 @@ class BeeperRoseGPT(nn.Module):
                 x = self._block_forward(blk, x)
         return self.norm(x)
 
-    def forward(self, idx: torch.Tensor) -> torch.Tensor:
+    def forward(self, idx: torch.Tensor, runtime_cfg: Optional[Dict[str, Any]] = None) -> torch.Tensor:
+        """
+        Forward pass with optional runtime pentachora influence.
+        If runtime_cfg is None, falls back to self.runtime_cfg set at init or via set_runtime_pentachora().
+        """
         h = self.backbone(idx)
+        cfg = self.runtime_cfg if runtime_cfg is None else {**self.runtime_cfg, **(runtime_cfg or {})}
+        h = self._apply_runtime_vertex_pull(h, cfg)
         return self.lm_head(h)
 
     # ---- Utilities ---------------------------------------------------------------------------------
@@ -245,7 +334,7 @@ class BeeperRoseGPT(nn.Module):
         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)."""
+        """Pool hidden states for Rose-related terms."""
         return h.mean(dim=1) if mode == "mean" else h[:, -1, :]
 
 
@@ -257,9 +346,7 @@ def prepare_model_for_state_dict(
 ) -> 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.
+    so we can load with strict=True. No-op if checkpoint lacks penta_* keys.
     """
     device = device or next(model.parameters()).device
     need = all(k in state_dict for k in ("penta_coarse", "penta_medium", "penta_fine"))
@@ -267,18 +354,15 @@ def prepare_model_for_state_dict(
         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.
+    def dims_ok(t: torch.Tensor, D: int) -> bool:
+        return t.ndim == 3 and t.size(1) == 5 and t.size(2) == D
+
+    D = model.token_emb.embedding_dim
+    if not (dims_ok(pc, D) and dims_ok(pt, D) and dims_ok(pm, D)):
         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)
+    model.ensure_pentachora(pc.size(0), pt.size(0), pm.size(0), dim=D, device=device)
 
 
 # --------------------------------- Generation -------------------------------------------------------
@@ -304,9 +388,10 @@ def generate(
     frequency_penalty: Optional[float] = None,
     device: Optional[torch.device] = None,
     detokenize: bool = True,
+    runtime_cfg: Optional[Dict[str, Any]] = None,   # <— NEW: pass-through to forward()
 ) -> str:
     """
-    Penalized nucleus sampling (same knobs as training script).
+    Penalized nucleus sampling with optional runtime pentachora influence.
     """
     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)
@@ -327,10 +412,10 @@ def generate(
             counts[t] += 1
 
     for _ in range(int(max_new_tokens)):
-        logits = model(x[:, -cfg["context"]:])
+        logits = model(x[:, -cfg["context"]:], runtime_cfg=runtime_cfg)
         logits = logits[:, -1, :]
 
-        # Repetition penalty (CTRL-like)
+        # Repetition penalty
         if repetition_penalty and repetition_penalty != 1.0:
             mask = counts > 0
             if mask.any():
@@ -338,7 +423,7 @@ def generate(
                 logits[:, mask][pos]  /= repetition_penalty
                 logits[:, mask][~pos] *= repetition_penalty
 
-        # Presence/frequency penalties (OpenAI-like)
+        # Presence/frequency penalties
         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)