Update app.py

app.py

@@ -1,29 +1,40 @@
-# app.py — CodVa-2 HF Space Inference
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import math
 import gradio as gr
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, HfApi
 from transformers import GPT2Tokenizer
 import json
 
-# ── Minimal inference build (copy the primitives from your training script) ────
+# ── Hardcoded config ──────────────────────────────────────────────────────────
+CFG = {
+    "vocab_size": 50304,
+    "d_model":    768,
+    "n_layers":   24,
+    "n_heads":    12,
+    "n_kv_heads": 3,
+    "ffn_dim":    1792,
+    "max_len":    2048,
+    "bus_points": [4, 9, 15, 21],
+}
+
+MODEL_REPO = "hugging-science/CodVa-2-session-001"
+
+# ── Primitives ────────────────────────────────────────────────────────────────
 
 class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
+    def __init__(self, dim, eps=1e-6):
         super().__init__()
         self.eps = eps
-        self.w   = nn.Parameter(torch.ones(dim))
+        self.w = nn.Parameter(torch.ones(dim))
     def forward(self, x):
         return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w
 
-def build_rope(head_dim: int, max_len: int, theta: float, device):
+def build_rope(head_dim, max_len, theta, device):
     freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
-    t     = torch.arange(max_len, device=device).float()
-    freqs = torch.outer(t, freqs)
-    return torch.cos(freqs), torch.sin(freqs)
+    t = torch.arange(max_len, device=device).float()
+    return torch.cos(torch.outer(t, freqs)), torch.sin(torch.outer(t, freqs))
 
 def apply_rope(x, cos, sin):
     B, H, L, D = x.shape
@@ -35,7 +46,7 @@ def apply_rope(x, cos, sin):
 def repeat_kv(x, n_rep):
     if n_rep == 1: return x
     B, H, L, D = x.shape
-    return x[:, :, None, :, :].expand(B, H, n_rep, L, D).reshape(B, H * n_rep, L, D)
+    return x[:, :, None, :, :].expand(B, H, n_rep, L, D).reshape(B, H*n_rep, L, D)
 
 class SwiGLU(nn.Module):
     def __init__(self, dim, hidden):
@@ -59,7 +70,6 @@ class GQA(nn.Module):
         self.wo = nn.Linear(n_heads * self.hd, dim, bias=False)
         self.q_norm = RMSNorm(self.hd)
         self.k_norm = RMSNorm(self.hd)
-
     def forward(self, x, cos, sin):
         B, L, _ = x.shape
         q = self.wq(x).view(B, L, self.n_heads, self.hd).transpose(1, 2)
@@ -79,272 +89,198 @@ class TransformerBlock(nn.Module):
         self.attn  = GQA(dim, n_heads, n_kv)
         self.ffn   = SwiGLU(dim, ffn_dim)
         self.res_scale = 1.0 / math.sqrt(2.0 * max(n_layers, 1))
-
     def forward(self, x, cos, sin):
         x = x + self.res_scale * self.attn(self.norm1(x), cos, sin)
         x = x + self.res_scale * self.ffn(self.norm2(x))
         return x
 
 class FiLMBridge(nn.Module):
-    def __init__(self, d_model):
+    def __init__(self, d):
         super().__init__()
-        self.conv    = nn.Conv1d(d_model, d_model, kernel_size=5, padding=2, groups=d_model)
-        self.to_film = nn.Linear(d_model, d_model * 2)
-
+        self.conv    = nn.Conv1d(d, d, 5, padding=2, groups=d)
+        self.to_film = nn.Linear(d, d * 2)
     def forward(self, h):
-        local       = self.conv(h.transpose(1, 2)).transpose(1, 2)
+        local = self.conv(h.transpose(1,2)).transpose(1,2)
         gamma, beta = self.to_film(local).chunk(2, dim=-1)
         return h * (1 + gamma) + beta
 
 class DepthTracker(nn.Module):
-    def __init__(self, d_model, d_state=32):
+    def __init__(self, d, d_state=32):
         super().__init__()
-        self.to_delta = nn.Linear(d_model, d_state)
-        self.to_out   = nn.Linear(d_state, d_model)
+        self.to_delta = nn.Linear(d, d_state)
+        self.to_out   = nn.Linear(d_state, d)
         self.decay    = nn.Parameter(torch.ones(d_state) * 0.9)
-
     def forward(self, x):
-        B, L, _  = x.shape
-        delta    = torch.tanh(self.to_delta(x))
-        decay    = torch.sigmoid(self.decay)
-        depth    = torch.zeros(B, delta.size(-1), device=x.device, dtype=x.dtype)
-        outs     = []
+        B, L, _ = x.shape
+        delta = torch.tanh(self.to_delta(x))
+        decay = torch.sigmoid(self.decay)
+        depth = torch.zeros(B, delta.size(-1), device=x.device, dtype=x.dtype)
+        outs = []
         for t in range(L):
             depth = depth * decay + delta[:, t]
             outs.append(depth)
         return self.to_out(torch.stack(outs, dim=1))
 
-class CNNBusContrib(nn.Module):
-    def __init__(self, d_model):
+class CNNBus(nn.Module):
+    def __init__(self, d):
         super().__init__()
-        self.conv_n = nn.Conv1d(d_model, d_model, 3, padding=1, groups=d_model)
-        self.conv_w = nn.Conv1d(d_model, d_model, 7, padding=3, groups=d_model)
-        self.proj   = nn.Linear(d_model * 2, d_model)
-
+        self.conv_n = nn.Conv1d(d, d, 3, padding=1, groups=d)
+        self.conv_w = nn.Conv1d(d, d, 7, padding=3, groups=d)
+        self.proj   = nn.Linear(d * 2, d)
     def forward(self, x):
         t = x.transpose(1, 2)
         n = F.silu(self.conv_n(t)).transpose(1, 2)
         w = F.silu(self.conv_w(t)).transpose(1, 2)
         return self.proj(torch.cat([n, w], dim=-1))
 
-class RNNBusContrib(nn.Module):
-    def __init__(self, d_model):
+class RNNBus(nn.Module):
+    def __init__(self, d):
         super().__init__()
-        self.gru  = nn.GRU(d_model, d_model // 2, num_layers=1, batch_first=True)
-        self.proj = nn.Linear(d_model // 2, d_model)
-
+        self.gru  = nn.GRU(d, d // 2, batch_first=True)
+        self.proj = nn.Linear(d // 2, d)
     def forward(self, x):
         out, _ = self.gru(x.float())
         return self.proj(out.to(x.dtype))
 
 class HighwayBus(nn.Module):
-    def __init__(self, d_model):
+    def __init__(self, d):
         super().__init__()
-        self.gate = nn.Linear(d_model * 2, d_model)
-
-    def forward(self, bus, contribution):
-        g = torch.sigmoid(self.gate(torch.cat([bus, contribution], dim=-1)))
-        return bus + g * contribution
+        self.gate = nn.Linear(d * 2, d)
+    def forward(self, bus, contrib):
+        g = torch.sigmoid(self.gate(torch.cat([bus, contrib], dim=-1)))
+        return bus + g * contrib
 
 class BusInjector(nn.Module):
     def __init__(self):
         super().__init__()
         self.gate = nn.Parameter(torch.zeros(1))
-
     def forward(self, h, bus):
         return h + torch.sigmoid(self.gate) * bus
 
-# ── CodVa-2 Inference Model ────
+# ── Model ─────────────────────────────────────────────────────────────────────
 
-class CodVa2Inference(nn.Module):
+class CodVa2(nn.Module):
     def __init__(self, cfg):
         super().__init__()
         self.cfg = cfg
-        
-        self.embed = nn.Embedding(cfg['vocab_size'], cfg['d_model'])
-        self.film = FiLMBridge(cfg['d_model'])
-        
-        self.blocks = nn.ModuleList([
-            TransformerBlock(cfg['d_model'], cfg['n_heads'], cfg['n_kv_heads'],
-                           cfg['ffn_dim'], n_layers=cfg['n_layers'])
+        D = cfg['d_model']
+        self.embed          = nn.Embedding(cfg['vocab_size'], D)
+        self.film           = FiLMBridge(D)
+        self.blocks         = nn.ModuleList([
+            TransformerBlock(D, cfg['n_heads'], cfg['n_kv_heads'],
+                             cfg['ffn_dim'], cfg['n_layers'])
             for _ in range(cfg['n_layers'])
         ])
-        
-        self.bus_points = cfg['bus_points']
-        self.highway = HighwayBus(cfg['d_model'])
-        self.injectors = nn.ModuleList([BusInjector() for _ in self.bus_points])
-        
-        self.depth_tracker = DepthTracker(cfg['d_model'])
-        self.cnn_1 = CNNBusContrib(cfg['d_model'])
-        self.rnn = RNNBusContrib(cfg['d_model'])
-        self.cnn_2 = CNNBusContrib(cfg['d_model'])
-        
-        self.norm_f = RMSNorm(cfg['d_model'])
-        self.lm_head = nn.Linear(cfg['d_model'], cfg['vocab_size'], bias=False)
+        self.bus_points     = cfg['bus_points']
+        self.highway        = HighwayBus(D)
+        self.injectors      = nn.ModuleList([BusInjector() for _ in self.bus_points])
+        self.depth_tracker  = DepthTracker(D)
+        self.cnn_1          = CNNBus(D)
+        self.rnn            = RNNBus(D)
+        self.cnn_2          = CNNBus(D)
         self.final_injector = BusInjector()
-        
-        self.register_buffer("cos", torch.zeros(cfg['max_len'], cfg['d_model'] // cfg['n_heads'] // 2))
-        self.register_buffer("sin", torch.zeros(cfg['max_len'], cfg['d_model'] // cfg['n_heads'] // 2))
-        
-    def _init_rope(self, device):
-        cos, sin = build_rope(self.cfg['d_model'] // self.cfg['n_heads'],
-                             self.cfg['max_len'], 500000.0, device)
+        self.norm_f         = RMSNorm(D)
+        self.lm_head        = nn.Linear(D, cfg['vocab_size'], bias=False)
+        hd = D // cfg['n_heads']
+        self.register_buffer("cos", torch.zeros(cfg['max_len'], hd // 2))
+        self.register_buffer("sin", torch.zeros(cfg['max_len'], hd // 2))
+
+    def _init_rope(self):
+        cos, sin = build_rope(
+            self.cfg['d_model'] // self.cfg['n_heads'],
+            self.cfg['max_len'], 500000.0,
+            self.embed.weight.device)
         self.cos.copy_(cos); self.sin.copy_(sin)
-    
+
     @torch.no_grad()
-    def generate(self, input_ids, max_new_tokens=100, temperature=0.7, top_p=0.9):
-        """Generate tokens autoregressively."""
-        device = self.embed.weight.device
-        
-        if self.cos.sum() == 0:
-            self._init_rope(device)
-        
-        if not isinstance(input_ids, torch.Tensor):
-            input_ids = torch.tensor([input_ids], device=device)
-        else:
-            input_ids = input_ids.to(device)
-        
-        if input_ids.dim() == 1:
-            input_ids = input_ids.unsqueeze(0)
-        
+    def generate(self, input_ids, max_new_tokens=128, temperature=0.8, top_p=0.92):
+        if self.cos.sum() == 0: self._init_rope()
+        ids = torch.tensor([input_ids]) if not isinstance(input_ids, torch.Tensor) \
+              else input_ids.unsqueeze(0)
+        bp = self.bus_points
         for _ in range(max_new_tokens):
-            L = input_ids.shape[1]
+            L = ids.shape[1]
             if L > self.cfg['max_len']:
-                input_ids = input_ids[:, -self.cfg['max_len']:]
-                L = self.cfg['max_len']
-            
-            # Forward pass
-            h = self.embed(input_ids)
-            h = self.film(h)
+                ids = ids[:, -self.cfg['max_len']:]
+                L = ids.shape[1]
+            h   = self.embed(ids)
+            h   = self.film(h)
             bus = torch.zeros_like(h)
-            
-            bp = self.bus_points
+            cos = self.cos[:L]; sin = self.sin[:L]
             for i, block in enumerate(self.blocks):
-                h = block(h, self.cos[:L], self.sin[:L])
-                if i == bp[0]:
-                    bus = self.highway(bus, self.depth_tracker(h))
-                    h = self.injectors[0](h, bus)
-                elif i == bp[1]:
-                    bus = self.highway(bus, self.cnn_1(h))
-                    h = self.injectors[1](h, bus)
-                elif i == bp[2]:
-                    h = self.injectors[2](h, bus)
-                elif i == bp[3]:
-                    bus = self.highway(bus, self.rnn(h))
-                    h = self.injectors[3](h, bus)
-            
+                h = block(h, cos, sin)
+                if   i == bp[0]: bus = self.highway(bus, self.depth_tracker(h)); h = self.injectors[0](h, bus)
+                elif i == bp[1]: bus = self.highway(bus, self.cnn_1(h));         h = self.injectors[1](h, bus)
+                elif i == bp[2]: h = self.injectors[2](h, bus)
+                elif i == bp[3]: bus = self.highway(bus, self.rnn(h));           h = self.injectors[3](h, bus)
             bus = self.highway(bus, self.cnn_2(h))
-            h = self.final_injector(h, bus)
-            h = self.norm_f(h)
-            
-            logits = self.lm_head(h[:, -1:, :])
-            logits = 30.0 * torch.tanh(logits / 30.0)
-            
-            # Sample
-            probs = F.softmax(logits / temperature, dim=-1)
-            sorted_probs, sorted_indices = torch.sort(probs, descending=True)
-            cumsum = torch.cumsum(sorted_probs, dim=-1)
-            mask = cumsum <= top_p
-            mask[..., 0] = True
-            filtered_probs = sorted_probs * mask.float()
-            filtered_probs /= filtered_probs.sum(dim=-1, keepdim=True) + 1e-9
-            
-            next_token = torch.multinomial(filtered_probs.squeeze(1), 1)
-            input_ids = torch.cat([input_ids, next_token], dim=1)
-        
-        return input_ids[0].tolist()
-
-# ── Load Model & Tokenizer ────
-
-MODEL_REPO = "hugging-science/CodVa-2-session-001"
-CHECKPOINT = "codva2_s001_step0001680.pt"
-
-print("Loading model...")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Download config
-config_path = hf_hub_download(MODEL_REPO, "config.json")
-with open(config_path) as f:
-    cfg = json.load(f)
-
-# Download checkpoint
-ckpt_path = hf_hub_download(MODEL_REPO, CHECKPOINT)
-
-# Build & load model
-model = CodVa2Inference(cfg).to(device)
-state = torch.load(ckpt_path, map_location=device)
-model.load_state_dict(state['model'], strict=False)
+            h   = self.final_injector(h, bus)
+            h   = self.norm_f(h)
+            logits = 30.0 * torch.tanh(self.lm_head(h[:, -1, :]) / 30.0)
+            probs  = F.softmax(logits / temperature, dim=-1)
+            sp, si = torch.sort(probs, descending=True)
+            mask   = (torch.cumsum(sp, dim=-1) - sp) < top_p
+            sp[~mask] = 0.0
+            sp /= sp.sum() + 1e-9
+            next_tok = si[torch.multinomial(sp, 1)]
+            ids = torch.cat([ids, next_tok.unsqueeze(0)], dim=1)
+            if next_tok.item() == 50256: break  # GPT2 EOS
+        return ids[0].tolist()
+
+# ── Load ──────────────────────────────────────────────────────────────────────
+
+print("Loading tokenizer...")
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+
+print("Finding latest checkpoint...")
+api    = HfApi()
+files  = list(api.list_repo_files(MODEL_REPO))
+ckpts  = sorted([f for f in files if f.endswith(".pt") and "step" in f],
+                key=lambda x: int(x.split("step")[-1].split(".")[0]))
+latest = ckpts[-1]
+print(f"Loading {latest}...")
+ckpt   = torch.load(hf_hub_download(MODEL_REPO, latest), map_location="cpu")
+
+model = CodVa2(CFG)
+model.load_state_dict(ckpt['model'], strict=False)
 model.eval()
+print(f"Ready! {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")
 
-# Load GPT2 tokenizer
-try:
-    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-except:
-    tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
-
-print(f"Model loaded: {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")
+# ── Gradio ────────────────────────────────────────────────────────────────────
 
-# ── Gradio Interface ────
-
-def generate_code(prompt: str, max_tokens: int = 100, temperature: float = 0.7, top_p: float = 0.9):
-    """Generate code from prompt."""
+def generate(prompt, max_tokens, temperature, top_p):
+    if not prompt.strip(): return "Enter a prompt."
     try:
-        # Tokenize
-        input_ids = tokenizer.encode(prompt, return_tensors="pt")[0]
-        
-        # Generate
-        output_ids = model.generate(input_ids.tolist(), max_new_tokens=max_tokens,
-                                   temperature=temperature, top_p=top_p)
-        
-        # Decode
-        output = tokenizer.decode(output_ids, skip_special_tokens=True)
-        
-        return output
+        ids = tokenizer.encode(prompt)[-256:]
+        out = model.generate(ids, int(max_tokens), float(temperature), float(top_p))
+        return tokenizer.decode(out, skip_special_tokens=True)
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error: {e}"
 
-# Gradio UI
-with gr.Blocks(title="CodVa-2") as demo:
-    gr.Markdown("# CodVa-2 — Deep-Narrow Code Model")
-    gr.Markdown("221M parameters, trained on 10B tokens of code+math")
-    
+with gr.Blocks(title="CodVa-2", theme=gr.themes.Monochrome()) as demo:
+    gr.Markdown("""
+    # 🍰 CodVa-2 — Code Model
+    **221M params** | 24L×768d | Highway Bus | GQA  
+    ⚠️ *CPU only — ~30-60s per generation. Early checkpoint, still training.*
+    """)
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="def fibonacci(n):",
-                lines=5,
-                value="def hello():"
-            )
-            
+            prompt = gr.Textbox(label="Prompt", lines=5,
+                                value="def fibonacci(n):\n    ")
             with gr.Row():
-                max_tokens = gr.Slider(10, 300, value=100, step=10, label="Max Tokens")
-                temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="Temperature")
-                top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-P")
-            
-            generate_btn = gr.Button("Generate", variant="primary")
-        
+                max_tok = gr.Slider(10, 200, 80,  step=10,   label="Max Tokens")
+                temp    = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="Temperature")
+                topp    = gr.Slider(0.1, 1.0, 0.92,step=0.05,label="Top-P")
+            btn = gr.Button("Generate 🚀", variant="primary")
         with gr.Column():
-            output = gr.Textbox(
-                label="Generated Code",
-                lines=10,
-                interactive=False
-            )
-    
-    generate_btn.click(
-        fn=generate_code,
-        inputs=[prompt, max_tokens, temperature, top_p],
-        outputs=output
-    )
-    
-    gr.Markdown("""
-    ## About CodVa-2
-    - **Architecture**: 24-layer deep-narrow transformer (768d) with GQA
-    - **Highway Bus**: FiLM bridge, DepthTracker, CNN experts, RNN expert
-    - **Training**: 10B mixed code+math tokens (StarCoder + FineWeb)
-    - **Stability**: Residual scaling, logit soft-cap, gradient clipping
-    - **Status**: Early checkpoint (step 1680) — expect improvement
-    """)
+            out = gr.Textbox(label="Output", lines=10, interactive=False)
+
+    btn.click(generate, [prompt, max_tok, temp, topp], out)
+
+    gr.Examples([
+        ["def fibonacci(n):\n    ", 100, 0.8, 0.92],
+        ["class Stack:\n    def __init__(self):\n        ", 120, 0.8, 0.92],
+        ["# binary search\ndef search(arr, target):\n    ", 100, 0.7, 0.9],
+    ], [prompt, max_tok, temp, topp])
 
-if __name__ == "__main__":
-    demo.launch(share=True)
+demo.launch()