fisrt

README.md

@@ -11,3 +11,11 @@ short_description: Local Self Attention Allows text to coherently self evolve
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# CNA — Simple Sampling (Strategy 1)
+
+This Space runs the **random position → argmax** update over a fixed-length sequence using your CNA checkpoint.
+
+## Quick Start
+
+1. In your Space repo, add your trained checkpoint as:

app.py

@@ -0,0 +1,263 @@
+# app.py
+import os, re, math, random
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import gradio as gr
+from transformers import AutoTokenizer
+
+# -----------------------------
+# Minimal CNA (inference-ready)
+# -----------------------------
+class AttnBlock(nn.Module):
+    def __init__(self, embed_dim, num_heads, expansion_factor):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim  = embed_dim // num_heads
+
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.QKV   = nn.Linear(embed_dim, embed_dim * 3)
+        self.Wo    = nn.Linear(embed_dim, embed_dim)
+
+        self.norm2 = nn.LayerNorm(embed_dim)
+        self.mlp   = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim * expansion_factor),
+            nn.GELU(),
+            nn.Linear(embed_dim * expansion_factor, embed_dim),
+        )
+
+        # match training's zero-init on residual branches
+        nn.init.zeros_(self.Wo.weight);   nn.init.zeros_(self.Wo.bias)
+        nn.init.zeros_(self.mlp[-1].weight); nn.init.zeros_(self.mlp[-1].bias)
+
+    def rope(self, Qh, Kh_seq, cos, sin):
+        Qe = Qh[..., 0::2]; Qo = Qh[..., 1::2]
+        ce = cos[..., 0::2]; se = sin[..., 0::2]
+        Qr_e = Qe * ce - Qo * se
+        Qr_o = Qe * se + Qo * ce
+        Qh2 = torch.empty_like(Qh); Qh2[..., 0::2] = Qr_e; Qh2[..., 1::2] = Qr_o
+
+        Ke = Kh_seq[..., 0::2]; Ko = Kh_seq[..., 1::2]
+        Kr_e = Ke * ce - Ko * se
+        Kr_o = Ke * se + Ko * ce
+        Kh2 = torch.empty_like(Kh_seq); Kh2[..., 0::2] = Kr_e; Kh2[..., 1::2] = Kr_o
+        return Qh2, Kh2
+
+    def forward(self, x, rope, radius):
+        h = self.norm1(x)
+        B, S, E = h.shape
+        cos, sin = rope
+        nh, hd = self.num_heads, self.head_dim
+
+        cos = cos.to(h.dtype).to(h.device).permute(0,2,1,3)  # [1,1,S,hd]
+        sin = sin.to(h.dtype).to(h.device).permute(0,2,1,3)
+
+        # local band mask
+        idx = torch.arange(S, device=h.device)
+        idx_dist = (idx.view(1, S) - idx.view(S, 1)).abs()
+        neg_inf = torch.finfo(h.dtype).min
+        mask = torch.full((S, S), neg_inf, dtype=h.dtype, device=h.device)
+        mask[idx_dist <= int(radius)] = 0
+        mask = mask.view(1, 1, S, S)
+
+        qkv = self.QKV(h)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        Qh = q.view(B,S,nh,hd).permute(0,2,1,3).contiguous()
+        Kh_seq = k.view(B,S,nh,hd).permute(0,2,1,3).contiguous()
+        Vh = v.view(B,S,nh,hd).permute(0,2,1,3).contiguous()
+
+        assert hd % 2 == 0, "rope needs even head_dim"
+        Qh, Kh_seq = self.rope(Qh, Kh_seq, cos, sin)
+        Kh = Kh_seq.permute(0,1,3,2).contiguous()
+
+        logits = (Qh @ Kh) * (hd ** -0.5)
+        attn = F.softmax(logits + mask, dim=-1) @ Vh
+        attn = attn.permute(0,2,1,3).contiguous().view(B,S,E)
+
+        x = x + self.Wo(attn)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+class CNA(nn.Module):
+    def __init__(self, embed_dim, num_heads, expansion_factor, num_blocks, radius, vocab_size):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.expansion_factor = expansion_factor
+        self.num_blocks = num_blocks
+        self.vocab_size = vocab_size
+        self.radius = radius
+        self.tok_emb = nn.Embedding(vocab_size, embed_dim)
+        self.blocks = nn.ModuleList([AttnBlock(embed_dim, num_heads, expansion_factor) for _ in range(num_blocks)])
+        self.proj = nn.Linear(embed_dim, vocab_size)
+
+    def _rope_seq(self, S, hd, device, dtype, base=10000.0):
+        pos = torch.arange(S, device=device, dtype=dtype)
+        half = hd // 2
+        idx = torch.arange(half, device=device, dtype=dtype)
+        inv = base ** (-idx / half)
+        ang = pos[:, None] * inv[None, :]
+        cos = ang.cos().unsqueeze(0).unsqueeze(2)
+        sin = ang.sin().unsqueeze(0).unsqueeze(2)
+        cos = torch.stack((cos, cos), dim=-1).reshape(1, S, 1, hd)
+        sin = torch.stack((sin, sin), dim=-1).reshape(1, S, 1, hd)
+        return cos, sin
+
+    def forward(self, x):
+        if x.dtype == torch.long and x.dim() == 2:
+            h = self.tok_emb(x)
+        else:
+            h = x
+        B, S, E = h.shape
+        hd = self.embed_dim // self.num_heads
+        cos, sin = self._rope_seq(S, hd, h.device, h.dtype)
+        for blk in self.blocks:
+            h = blk(h, rope=(cos, sin), radius=self.radius)
+        return self.proj(h)
+
+# -----------------------------
+# Helpers (trimmed to Strategy 1)
+# -----------------------------
+def infer_expansion_factor_from_state(state, embed_dim):
+    for key in ("blocks.0.mlp.0.weight", "blocks.0.mlp.2.weight"):
+        if key in state:
+            W = state[key]
+            if key.endswith("0.weight"):
+                return int(W.shape[0] // embed_dim)
+            else:
+                return int(W.shape[1] // embed_dim)
+    return 4
+
+@torch.no_grad()
+def decode(ids, tokenizer, max_chars=220):
+    s = tokenizer.decode(ids.tolist(), skip_special_tokens=True)
+    s = s.replace("\n", " ")
+    return s[:max_chars] + ("…" if len(s) > max_chars else "")
+
+@torch.no_grad()
+def model_logits(model, x):
+    return model(x)
+
+# -----------------------------
+# Load checkpoint & build model
+# -----------------------------
+def load_model(ckpt_path: str):
+    if not os.path.exists(ckpt_path):
+        raise FileNotFoundError(
+            f"Checkpoint not found at {ckpt_path}. "
+            "Upload ckpt_latest.pt to the repo root or set the correct path."
+        )
+    payload = torch.load(ckpt_path, map_location="cpu")
+    state = payload["model"]
+    cfg = payload.get("config", {}) or {}
+
+    # Carry over config (robust fallbacks)
+    embed_dim        = cfg.get("embed_dim")
+    num_heads        = cfg.get("num_heads")
+    num_blocks       = cfg.get("num_blocks")
+    radius           = cfg.get("radius")
+    expansion_factor = cfg.get("expansion_factor")
+
+    if embed_dim is None:  embed_dim = state["tok_emb.weight"].shape[1]
+    if num_blocks is None:
+        block_idxs = [int(m.group(1)) for k in state.keys() for m in [re.match(r"blocks\.(\d+)\.", k)] if m]
+        num_blocks = max(block_idxs) + 1 if block_idxs else 1
+    if num_heads is None: num_heads = 8
+    if radius is None: radius = 16
+    if expansion_factor is None:
+        expansion_factor = infer_expansion_factor_from_state(state, embed_dim)
+    else:
+        expansion_factor = int(expansion_factor)
+
+    tokenizer_name = payload.get("tokenizer_name", "gpt2")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.model_max_length = 1_000_000_000
+    vocab_size = tokenizer.vocab_size
+
+    model = CNA(
+        int(embed_dim), int(num_heads), int(expansion_factor),
+        int(num_blocks), int(radius), int(vocab_size)
+    )
+
+    # Load weights (tolerate proj head size diff)
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    if any(k.startswith("proj.") for k in missing):
+        with torch.no_grad():
+            nn.init.normal_(model.proj.weight, std=0.02)
+            nn.init.zeros_(model.proj.bias)
+    else:
+        model.load_state_dict(state, strict=True)
+
+    model.eval()
+    return model, tokenizer, int(radius)
+
+# -----------------------------
+# Simplest sampling: Strategy 1
+# -----------------------------
+@torch.no_grad()
+def strategy_random_argmax(model, tokenizer, seqlen=100, steps=200, snap_every=20, seed=0, max_chars=220):
+    random.seed(seed); torch.manual_seed(seed)
+    V = tokenizer.vocab_size
+    x = torch.randint(0, V, (1, seqlen))
+    snaps = [(0, decode(x[0].cpu(), tokenizer, max_chars))]
+    for t in range(1, steps + 1):
+        pos = int(torch.randint(0, seqlen, (1,)))
+        logits_pos = model_logits(model, x)[0, pos]  # [V]
+        x[0, pos] = int(torch.argmax(logits_pos).item())
+        if (t % snap_every == 0) or (t == steps):
+            snaps.append((t, decode(x[0].cpu(), tokenizer, max_chars)))
+    return snaps
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+DEFAULT_CKPT = os.environ.get("CKPT_PATH", "ckpt_latest.pt")
+
+model_cache = {"model": None, "tokenizer": None, "radius": None, "ckpt": None}
+def ensure_model(ckpt_path):
+    if model_cache["model"] is None or model_cache["ckpt"] != ckpt_path:
+        m, tok, rad = load_model(ckpt_path)
+        model_cache.update({"model": m, "tokenizer": tok, "radius": rad, "ckpt": ckpt_path})
+
+def run_demo(ckpt_path, seqlen, steps, snap_every, seed, max_chars):
+    ensure_model(ckpt_path or DEFAULT_CKPT)
+    snaps = strategy_random_argmax(
+        model_cache["model"], model_cache["tokenizer"],
+        seqlen=seqlen, steps=steps, snap_every=snap_every,
+        seed=seed, max_chars=max_chars
+    )
+    # Pretty print log
+    log = "\n".join([f"t={t:>3}: {txt}" for (t, txt) in snaps])
+    final_text = snaps[-1][1] if snaps else ""
+    return log, final_text
+
+with gr.Blocks(title="CNA — Simple Sampling (Random Position • Argmax)") as demo:
+    gr.Markdown(
+        """
+        # CNA — Simple Sampling (Strategy 1)
+        This Space loads your checkpoint and runs the **random position → argmax** update for a fixed-length sequence.
+        - Put your checkpoint at `ckpt_latest.pt` (repo root), or set a custom path below.
+        """
+    )
+    with gr.Row():
+        ckpt = gr.Textbox(value=DEFAULT_CKPT, label="Checkpoint path", placeholder="ckpt_latest.pt")
+    with gr.Row():
+        seqlen = gr.Slider(10, 512, value=100, step=1, label="Sequence length (S)")
+        steps = gr.Slider(10, 1000, value=200, step=1, label="Steps")
+        snap_every = gr.Slider(1, 200, value=20, step=1, label="Snapshot every N steps")
+    with gr.Row():
+        seed = gr.Slider(0, 10_000, value=0, step=1, label="Seed")
+        max_chars = gr.Slider(32, 1000, value=220, step=1, label="Max chars per snapshot")
+    run_btn = gr.Button("Run")
+    with gr.Row():
+        log_out = gr.Textbox(lines=18, label="Snapshots")
+    final_out = gr.Textbox(lines=6, label="Final text (last snapshot)")
+
+    run_btn.click(run_demo, [ckpt, seqlen, steps, snap_every, seed, max_chars], [log_out, final_out])
+
+demo.queue(concurrency_count=1).launch()