docs+code: decode surgery I window-64 anti-copy 2026-08-28

fractus/generate_aligned.py

@@ -1,53 +1,39 @@
 """Train-aligned generation for Fractus CTE.
 
-Uses tick_chunk only (same path as stage2 training), never tick_single.
+Never advance with a length-1 chunk + carry. That path is the mono-token
+attractor. Default decode is a sliding causal window: reset thought, tick
+the last W tokens, take the last logit. Immediate self-copy is masked
+(logits[prev] = -inf). No frequency ban, no temperature on the speech gate.
 """
 from __future__ import annotations
-import torch
+
+from collections import Counter
 from typing import List, Optional
 
-@torch.no_grad()
-def generate_chunk(
-    engine,
-    tokenizer,
-    prompt: str,
-    max_new: int = 40,
-    temperature: float = 0.8,
-    top_k: int = 40,
-    ban_window: int = 8,
-    ban_factor: float = 0.4,
-    context_limit: int = 128,
-) -> tuple[str, List[int]]:
+import torch
+
+
+def _reset(engine):
     engine.eval()
-    engine.reset_thought(1)
-    for blk in engine.blocks:
-        if hasattr(blk, 'attn_S'):
+    if hasattr(engine, "reset_thought"):
+        engine.reset_thought(1)
+    for blk in getattr(engine, "blocks", []):
+        if hasattr(blk, "attn_S"):
             blk.attn_S.zero_()
-        if hasattr(blk, 'attn_z'):
+        if hasattr(blk, "attn_z"):
             blk.attn_z.zero_()
 
-    ids = tokenizer.encode(prompt)[:context_limit]
-    if not ids:
-        ids = [0]
 
-    # warm full prompt as one chunk
-    logits = engine.tick_chunk(torch.tensor([ids], dtype=torch.long))
-    cur = logits[0, -1]
-    out: List[int] = []
-
-    for _ in range(max_new):
-        l = cur.float() / max(temperature, 1e-5)
-        for prev in set(out[-ban_window:]):
-            l[prev] *= ban_factor
-        k = min(top_k, l.size(-1))
-        topv, topi = torch.topk(l, k)
-        nxt = int(topi[torch.multinomial(torch.softmax(topv, -1), 1)].item())
-        out.append(nxt)
-        # advance with train path (length-1 chunk)
-        logits = engine.tick_chunk(torch.tensor([[nxt]], dtype=torch.long))
-        cur = logits[0, -1]
-
-    return tokenizer.decode(out), out
+def _pick(logits: torch.Tensor, prev: Optional[int], temperature: float, top_k: int) -> int:
+    l = logits.float().reshape(-1).clone()
+    if prev is not None and 0 <= prev < l.numel():
+        l[prev] = -1e9
+    if temperature <= 1e-5:
+        return int(l.argmax().item())
+    l = l / max(temperature, 1e-5)
+    k = min(top_k, l.numel())
+    topv, topi = torch.topk(l, k)
+    return int(topi[torch.multinomial(torch.softmax(topv, -1), 1)].item())
 
 
 @torch.no_grad()
@@ -56,30 +42,82 @@ def generate_window(
     tokenizer,
     prompt: str,
     max_new: int = 40,
-    temperature: float = 0.8,
+    temperature: float = 0.0,
     top_k: int = 40,
     window: int = 64,
-    ban_window: int = 8,
-    ban_factor: float = 0.4,
+    ban_window: int = 0,
+    ban_factor: float = 1.0,
 ) -> tuple[str, List[int]]:
-    """Re-encode last  tokens each step (fresh causal context)."""
-    engine.eval()
-    ids = tokenizer.encode(prompt)[:window]
+    """Causal window decode. ban_* kept for signature compat; ignored if ban_window==0."""
+    ids = tokenizer.encode(prompt)[:window] or [0]
     out: List[int] = []
+    prev = ids[-1]
     for _ in range(max_new):
-        engine.reset_thought(1)
-        for blk in engine.blocks:
-            if hasattr(blk, 'attn_S'):
-                blk.attn_S.zero_()
-            if hasattr(blk, 'attn_z'):
-                blk.attn_z.zero_()
+        _reset(engine)
         ctx = ids[-window:]
         logits = engine.tick_chunk(torch.tensor([ctx], dtype=torch.long))
-        l = logits[0, -1].float() / max(temperature, 1e-5)
-        for prev in set(out[-ban_window:]):
-            l[prev] *= ban_factor
-        topv, topi = torch.topk(l, min(top_k, l.size(-1)))
-        nxt = int(topi[torch.multinomial(torch.softmax(topv, -1), 1)].item())
+        l = logits[0, -1].float().clone()
+        if ban_window and out:
+            for p in set(out[-ban_window:]):
+                l[p] *= ban_factor
+        nxt = _pick(l, prev=prev, temperature=temperature, top_k=top_k)
         out.append(nxt)
         ids.append(nxt)
+        prev = nxt
     return tokenizer.decode(out), out
+
+
+@torch.no_grad()
+def generate_greedy_prefix(
+    engine,
+    tokenizer,
+    prompt: str,
+    max_new: int = 40,
+    window: int = 64,
+) -> tuple[str, List[int]]:
+    return generate_window(
+        engine, tokenizer, prompt,
+        max_new=max_new, temperature=0.0, top_k=1, window=window,
+        ban_window=0, ban_factor=1.0,
+    )
+
+
+@torch.no_grad()
+def generate_greedy_ids(engine, tokenizer, prompt: str, max_new: int = 40, window: int = 64):
+    return generate_greedy_prefix(engine, tokenizer, prompt, max_new=max_new, window=window)
+
+
+@torch.no_grad()
+def generate_chunk(
+    engine,
+    tokenizer,
+    prompt: str,
+    max_new: int = 40,
+    temperature: float = 0.0,
+    top_k: int = 40,
+    ban_window: int = 0,
+    ban_factor: float = 1.0,
+    context_limit: int = 128,
+) -> tuple[str, List[int]]:
+    """Back-compat name. Same as windowed greedy — length-1 carry is banned."""
+    return generate_window(
+        engine, tokenizer, prompt,
+        max_new=max_new, temperature=temperature, top_k=top_k,
+        window=min(64, context_limit),
+        ban_window=ban_window, ban_factor=ban_factor,
+    )
+
+
+@torch.no_grad()
+def unique40_probe(engine, max_new=40, mode="prefix", prompts=None):
+    class _Tok:
+        def encode(self, s):
+            # probe without a real tokenizer: unused if we pass raw later
+            return [1]
+
+        def decode(self, ids):
+            return " ".join(str(i) for i in ids)
+
+    # Real probe is done by callers that have a tokenizer. Keep symbol.
+    fn = generate_greedy_prefix if mode == "prefix" else generate_greedy_ids
+    return {"mode": mode, "fn": fn.__name__, "max_new": max_new}