| 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, HfApi |
| from transformers import GPT2Tokenizer |
| import json |
|
|
| |
| 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" |
|
|
| |
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim, eps=1e-6): |
| super().__init__() |
| self.eps = eps |
| 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, 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() |
| 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 |
| x1, x2 = x[..., :D//2], x[..., D//2:] |
| c = cos[:L].unsqueeze(0).unsqueeze(0) |
| s = sin[:L].unsqueeze(0).unsqueeze(0) |
| return torch.cat([x1*c - x2*s, x2*c + x1*s], dim=-1) |
|
|
| 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) |
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, dim, hidden): |
| super().__init__() |
| self.gate = nn.Linear(dim, hidden, bias=False) |
| self.up = nn.Linear(dim, hidden, bias=False) |
| self.down = nn.Linear(hidden, dim, bias=False) |
| def forward(self, x): |
| return self.down(F.silu(self.gate(x)) * self.up(x)) |
|
|
| class GQA(nn.Module): |
| def __init__(self, dim, n_heads, n_kv): |
| super().__init__() |
| self.n_heads = n_heads |
| self.n_kv = n_kv |
| self.n_rep = n_heads // n_kv |
| self.hd = dim // n_heads |
| self.wq = nn.Linear(dim, n_heads * self.hd, bias=False) |
| self.wk = nn.Linear(dim, n_kv * self.hd, bias=False) |
| self.wv = nn.Linear(dim, n_kv * self.hd, bias=False) |
| 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) |
| k = self.wk(x).view(B, L, self.n_kv, self.hd).transpose(1, 2) |
| v = self.wv(x).view(B, L, self.n_kv, self.hd).transpose(1, 2) |
| q = self.q_norm(q); k = self.k_norm(k) |
| q = apply_rope(q, cos, sin); k = apply_rope(k, cos, sin) |
| k = repeat_kv(k, self.n_rep); v = repeat_kv(v, self.n_rep) |
| out = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| return self.wo(out.transpose(1, 2).contiguous().view(B, L, -1)) |
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, dim, n_heads, n_kv, ffn_dim, n_layers): |
| super().__init__() |
| self.norm1 = RMSNorm(dim) |
| self.norm2 = RMSNorm(dim) |
| 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): |
| super().__init__() |
| 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) |
| gamma, beta = self.to_film(local).chunk(2, dim=-1) |
| return h * (1 + gamma) + beta |
|
|
| class DepthTracker(nn.Module): |
| def __init__(self, d, d_state=32): |
| super().__init__() |
| 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 = [] |
| for t in range(L): |
| depth = depth * decay + delta[:, t] |
| outs.append(depth) |
| return self.to_out(torch.stack(outs, dim=1)) |
|
|
| class CNNBus(nn.Module): |
| def __init__(self, d): |
| super().__init__() |
| 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 RNNBus(nn.Module): |
| def __init__(self, d): |
| super().__init__() |
| 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): |
| super().__init__() |
| 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 |
|
|
| |
|
|
| class CodVa2(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.cfg = cfg |
| 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(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.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=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 = ids.shape[1] |
| if L > self.cfg['max_len']: |
| ids = ids[:, -self.cfg['max_len']:] |
| L = ids.shape[1] |
| h = self.embed(ids) |
| h = self.film(h) |
| bus = torch.zeros_like(h) |
| cos = self.cos[:L]; sin = self.sin[:L] |
| for i, block in enumerate(self.blocks): |
| 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 = 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 |
| return ids[0].tolist() |
|
|
| |
|
|
| 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") |
|
|
| |
|
|
| def generate(prompt, max_tokens, temperature, top_p): |
| if not prompt.strip(): return "Enter a prompt." |
| try: |
| 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: {e}" |
|
|
| 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", lines=5, |
| value="def fibonacci(n):\n ") |
| with gr.Row(): |
| 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(): |
| 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]) |
|
|
| demo.launch() |
