fix generate_window kwargs force_second carry 2026-08-29b

fractus/generate_aligned.py

@@ -1,18 +1,8 @@
-"""Train-aligned generation for Fractus CTE.
-
-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.
-"""
+"""Train-aligned generation for Fractus CTE. No frequency ban / scramble."""
 from __future__ import annotations
-
-from collections import Counter
 from typing import List, Optional
-
 import torch
 
-
 def _reset(engine):
     engine.eval()
     if hasattr(engine, "reset_thought"):
@@ -23,7 +13,6 @@ def _reset(engine):
         if hasattr(blk, "attn_z"):
             blk.attn_z.zero_()
 
-
 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():
@@ -31,11 +20,10 @@ def _pick(logits: torch.Tensor, prev: Optional[int], temperature: float, top_k:
     if temperature <= 1e-5:
         return int(l.argmax().item())
     l = l / max(temperature, 1e-5)
-    k = min(top_k, l.numel())
+    k = min(max(1, top_k), l.numel())
     topv, topi = torch.topk(l, k)
     return int(topi[torch.multinomial(torch.softmax(topv, -1), 1)].item())
 
-
 @torch.no_grad()
 def generate_window(
     engine,
@@ -44,80 +32,60 @@ def generate_window(
     max_new: int = 40,
     temperature: float = 0.0,
     top_k: int = 40,
-    window: int = 64,
-    ban_window: int = 0,
-    ban_factor: float = 1.0,
+    window: int = 128,
+    force_second: bool = False,
+    carry: bool = False,
 ) -> tuple[str, List[int]]:
-    """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):
+    if carry:
         _reset(engine)
-        ctx = ids[-window:]
-        logits = engine.tick_chunk(torch.tensor([ctx], dtype=torch.long))
-        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
+        logits = engine.tick_chunk(torch.tensor([ids], dtype=torch.long))
+        cur = logits[0, -1]
+        for i in range(max_new):
+            if force_second and i == 0:
+                l = cur.float().clone()
+                l[int(l.argmax())] = -1e9
+                nxt = _pick(l, prev=prev, temperature=temperature, top_k=top_k)
+            else:
+                nxt = _pick(cur, prev=prev, temperature=temperature, top_k=top_k)
+            out.append(nxt)
+            ids.append(nxt)
+            logits = engine.tick_chunk(torch.tensor([[nxt]], dtype=torch.long))
+            cur = logits[0, -1]
+            prev = nxt
+    else:
+        for i in range(max_new):
+            _reset(engine)
+            ctx = ids[-window:]
+            logits = engine.tick_chunk(torch.tensor([ctx], dtype=torch.long))
+            cur = logits[0, -1]
+            if force_second and i == 0:
+                l = cur.float().clone()
+                l[int(l.argmax())] = -1e9
+                nxt = _pick(l, prev=prev, temperature=temperature, top_k=top_k)
+            else:
+                nxt = _pick(cur, 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,
-    )
-
+def generate_greedy_prefix(engine, tokenizer, prompt: str, max_new: int = 40, window: int = 128):
+    return generate_window(engine, tokenizer, prompt, max_new=max_new, temperature=0.0, top_k=1, window=window, carry=False)
 
 @torch.no_grad()
-def generate_greedy_ids(engine, tokenizer, prompt: str, max_new: int = 40, window: int = 64):
+def generate_greedy_ids(engine, tokenizer, prompt: str, max_new: int = 40, window: int = 128):
     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,
-    )
-
+def generate_chunk(engine, tokenizer, prompt: str, max_new: int = 40, **kw):
+    return generate_window(engine, tokenizer, prompt, max_new=max_new, window=128, **kw)
 
 @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)
+    return {"mode": mode, "max_new": max_new}
 
-    # 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}
+# space-sync 2026-08-29b force_second carry