| |
|
|
| import os |
| import math |
| import time |
|
|
| |
| |
| |
| |
| os.environ["OMP_NUM_THREADS"] = "2" |
| os.environ["MKL_NUM_THREADS"] = "2" |
| os.environ["OPENBLAS_NUM_THREADS"] = "1" |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import gradio as gr |
| from dataclasses import dataclass |
| from typing import Tuple, Generator |
| from tokenizers import Tokenizer |
| from huggingface_hub import hf_hub_download, login, HfApi |
|
|
| torch.set_num_threads(2) |
| torch.set_num_interop_threads(1) |
|
|
| print(f"[cpu] 2 vCPUs | torch threads: {torch.get_num_threads()}") |
|
|
| |
| |
| |
| HF_TOKEN = os.environ.get("HF_TOKEN", "") |
| HF_DATASET_REPO = "Bc-AI/nova1_data" |
| HF_MODEL_REPO = os.environ.get("MODEL_REPO", "hugging-science/CodVa-2-session-002") |
|
|
| if HF_TOKEN: |
| login(token=HF_TOKEN) |
|
|
| |
| |
| |
| @dataclass |
| class Config: |
| vocab_size: int = 50304 |
| d_model: int = 896 |
| n_layers: int = 18 |
| n_heads: int = 14 |
| n_kv_heads: int = 2 |
| max_len: int = 2048 |
| rope_theta: float = 500_000.0 |
| window_size: int = 512 |
| pattern_mult: float = 2.5 |
| reason_mult: float = 0.75 |
| reason_depth: int = 2 |
| gate_hidden: int = 64 |
| gate_init: float = 0.0 |
| diff_lambda_init: float = 0.8 |
| tie_embeddings: bool = True |
|
|
| @property |
| def head_dim(self): |
| return self.d_model // self.n_heads |
|
|
| @property |
| def pattern_dim(self): |
| return ((int(self.d_model * self.pattern_mult) + 255) // 256) * 256 |
|
|
| @property |
| def reason_dim(self): |
| return ((int(self.d_model * self.reason_mult) + 63) // 64) * 64 |
|
|
| |
| |
| |
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-6): |
| super().__init__() |
| self.eps = eps |
| self.w = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| x32 = x.float() |
| return (x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps) * self.w).to(x.dtype) |
|
|
|
|
| def precompute_rope(head_dim, max_len, theta, device): |
| inv_freq = 1.0 / (theta ** ( |
| torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim |
| )) |
| pos = torch.arange(max_len, device=device, dtype=torch.float32) |
| freqs = torch.outer(pos, inv_freq) |
| return freqs.cos(), freqs.sin() |
|
|
|
|
| def apply_rope(x, cos, sin): |
| B, H, L, D = x.shape |
| h = D // 2 |
| c = cos[:L].unsqueeze(0).unsqueeze(0) |
| s = sin[:L].unsqueeze(0).unsqueeze(0) |
| return torch.cat([ |
| x[..., :h] * c - x[..., h:] * s, |
| x[..., h:] * c + x[..., :h] * s, |
| ], dim=-1) |
|
|
|
|
| def repeat_kv(x, n_rep): |
| if n_rep == 1: return x |
| B, H, L, D = x.shape |
| return x.unsqueeze(2).expand(B, H, n_rep, L, D).reshape(B, H * n_rep, L, D) |
|
|
|
|
| class DifferentialAttention(nn.Module): |
| def __init__(self, cfg, layer_idx, local=False): |
| super().__init__() |
| self.n_pairs = cfg.n_heads // 2 |
| self.n_kv_pairs = max(1, cfg.n_kv_heads // 2) |
| self.n_rep = self.n_pairs // self.n_kv_pairs |
| self.head_dim = cfg.head_dim |
| self.local = local |
| self.window = cfg.window_size |
| d = cfg.d_model |
|
|
| self.wq = nn.Linear(d, 2 * self.n_pairs * self.head_dim, bias=False) |
| self.wk = nn.Linear(d, 2 * self.n_kv_pairs * self.head_dim, bias=False) |
| self.wv = nn.Linear(d, self.n_kv_pairs * self.head_dim, bias=False) |
| self.wo = nn.Linear( self.n_pairs * self.head_dim, d, bias=False) |
| self.q_norm = RMSNorm(self.head_dim) |
| self.k_norm = RMSNorm(self.head_dim) |
| self.lambda1 = nn.Parameter(torch.tensor(0.0)) |
| self.lambda2 = nn.Parameter(torch.tensor(0.0)) |
| self.out_norm = RMSNorm(self.head_dim) |
|
|
| def _window_mask(self, L, device): |
| idx = torch.arange(L, device=device) |
| dist = idx.unsqueeze(0) - idx.unsqueeze(1) |
| mask = (dist > 0) | (dist < -self.window) |
| return mask.float().masked_fill(mask, float('-inf')) |
|
|
| def forward(self, x, cos, sin): |
| B, L, _ = x.shape |
| hd = self.head_dim |
|
|
| q_all = self.wq(x).view(B, L, self.n_pairs, 2, hd).transpose(1, 2) |
| k_all = self.wk(x).view(B, L, self.n_kv_pairs, 2, hd).transpose(1, 2) |
| v = self.wv(x).view(B, L, self.n_kv_pairs, hd).transpose(1, 2) |
| q1, q2 = q_all[..., 0, :], q_all[..., 1, :] |
| k1, k2 = k_all[..., 0, :], k_all[..., 1, :] |
|
|
| q1 = self.q_norm(q1); q2 = self.q_norm(q2) |
| k1 = self.k_norm(k1); k2 = self.k_norm(k2) |
| q1 = apply_rope(q1, cos, sin); q2 = apply_rope(q2, cos, sin) |
| k1 = apply_rope(k1, cos, sin); k2 = apply_rope(k2, cos, sin) |
| k1 = repeat_kv(k1, self.n_rep); k2 = repeat_kv(k2, self.n_rep) |
| v = repeat_kv(v, self.n_rep) |
|
|
| scale = 1.0 / math.sqrt(hd) |
| mask = (self._window_mask(L, x.device) if self.local |
| else torch.zeros(L, L, device=x.device).masked_fill( |
| ~torch.ones(L, L, device=x.device, dtype=torch.bool).tril(), |
| float('-inf'))) |
|
|
| |
| a1 = F.softmax(torch.matmul(q1, k1.transpose(-2, -1)) * scale + mask, dim=-1) |
| a2 = F.softmax(torch.matmul(q2, k2.transpose(-2, -1)) * scale + mask, dim=-1) |
| lam = torch.exp(self.lambda1) - torch.exp(self.lambda2) + 0.5 |
| out = torch.matmul(a1 - lam * a2, v) |
| out = self.out_norm(out).transpose(1, 2).contiguous().view(B, L, -1) |
| return self.wo(out) |
|
|
|
|
| class DualPathFFN(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| d, pd, rd = cfg.d_model, cfg.pattern_dim, cfg.reason_dim |
| self.pat_gate = nn.Linear(d, pd, bias=False) |
| self.pat_up = nn.Linear(d, pd, bias=False) |
| self.pat_down = nn.Linear(pd, d, bias=False) |
| layers = [nn.Linear(d, rd, bias=False), nn.SiLU()] |
| for _ in range(cfg.reason_depth - 1): |
| layers += [nn.Linear(rd, rd, bias=False), nn.SiLU()] |
| layers.append(nn.Linear(rd, d, bias=False)) |
| self.reason = nn.Sequential(*layers) |
| self.merge = nn.Parameter(torch.zeros(d)) |
|
|
| def forward(self, x): |
| pat = self.pat_down(F.silu(self.pat_gate(x)) * self.pat_up(x)) |
| w = torch.sigmoid(self.merge) |
| return w * pat + (1.0 - w) * self.reason(x) |
|
|
|
|
| class TokenImportanceGate(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(cfg.d_model, cfg.gate_hidden, bias=True), |
| nn.SiLU(), |
| nn.Linear(cfg.gate_hidden, 1, bias=True), |
| ) |
| def forward(self, x): |
| return x * (0.5 + torch.sigmoid(self.net(x))) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg, layer_idx): |
| super().__init__() |
| self.norm1 = RMSNorm(cfg.d_model) |
| self.norm2 = RMSNorm(cfg.d_model) |
| self.attn = DifferentialAttention(cfg, layer_idx, local=(layer_idx % 2 == 0)) |
| self.ffn = DualPathFFN(cfg) |
|
|
| def forward(self, x, cos, sin): |
| x = x + self.attn(self.norm1(x), cos, sin) |
| x = x + self.ffn(self.norm2(x)) |
| return x |
|
|
|
|
| class CodVa2(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.cfg = cfg |
| self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model) |
| self.importance = TokenImportanceGate(cfg) |
| self.blocks = nn.ModuleList([Block(cfg, i) for i in range(cfg.n_layers)]) |
| self.final_norm = RMSNorm(cfg.d_model) |
| self.register_buffer("rope_cos", torch.zeros(cfg.max_len, cfg.head_dim // 2)) |
| self.register_buffer("rope_sin", torch.zeros(cfg.max_len, cfg.head_dim // 2)) |
| self._rope_ready = False |
|
|
| def _init_rope(self, device): |
| c, s = precompute_rope( |
| self.cfg.head_dim, self.cfg.max_len, self.cfg.rope_theta, device |
| ) |
| self.rope_cos.copy_(c) |
| self.rope_sin.copy_(s) |
| self._rope_ready = True |
|
|
| def forward(self, tokens): |
| B, L = tokens.shape |
| device = tokens.device |
| if not self._rope_ready: |
| self._init_rope(device) |
| cos = self.rope_cos[:L] |
| sin = self.rope_sin[:L] |
| x = self.importance(self.embed(tokens)) |
| for block in self.blocks: |
| x = block(x, cos, sin) |
| return F.linear(self.final_norm(x), self.embed.weight) |
|
|
| |
| |
| |
| print("[init] loading tokenizer...") |
| tok_path = hf_hub_download( |
| repo_id=HF_DATASET_REPO, |
| filename="nova_tokenizer.json", |
| repo_type="dataset", |
| token=HF_TOKEN or None, |
| ) |
| tokenizer = Tokenizer.from_file(tok_path) |
|
|
| DOMAIN_TOKENS = { |
| "Code": "<|domain_code|>", |
| "Math": "<|domain_math|>", |
| "General": "<|domain_general|>", |
| "Reasoning": "<|domain_reasoning|>", |
| } |
|
|
| EOS_ID = tokenizer.token_to_id("<|endoftext|>") or tokenizer.token_to_id("</s>") or -1 |
| print(f"[init] tokenizer | vocab={tokenizer.get_vocab_size()} | eos={EOS_ID}") |
|
|
| print("[init] loading model...") |
| cfg = Config() |
| cfg.vocab_size = (tokenizer.get_vocab_size() + 63) // 64 * 64 |
| model = CodVa2(cfg) |
|
|
| api = HfApi() |
| files = list(api.list_repo_files( |
| repo_id=HF_MODEL_REPO, repo_type="model", token=HF_TOKEN or None |
| )) |
| finals = sorted([f for f in files if "final" in f and f.endswith(".safetensors")]) |
| ckpts = sorted( |
| [f for f in files if "step" in f and f.endswith(".pt")], |
| key=lambda x: int(x.split("step")[-1].split(".")[0]) |
| ) |
|
|
| if finals: |
| print(f"[init] loading final: {finals[-1]}") |
| from safetensors.torch import load_file |
| wpath = hf_hub_download(HF_MODEL_REPO, finals[-1], |
| repo_type="model", token=HF_TOKEN or None) |
| model.load_state_dict(load_file(wpath, device="cpu"), strict=True) |
|
|
| elif ckpts: |
| print(f"[init] loading checkpoint: {ckpts[-1]}") |
| wpath = hf_hub_download(HF_MODEL_REPO, ckpts[-1], |
| repo_type="model", token=HF_TOKEN or None) |
| ckpt = torch.load(wpath, map_location="cpu", weights_only=False) |
| model.load_state_dict(ckpt["model"], strict=True) |
|
|
| else: |
| print("[init] WARNING: no weights found") |
|
|
| model.eval() |
| n_params = sum(p.numel() for p in model.parameters()) |
| print(f"[init] ready | {n_params/1e6:.1f}M params | 2 vCPUs (free tier)") |
|
|
| |
| |
| |
| @torch.no_grad() |
| def generate_stream( |
| prompt: str, |
| domain: str, |
| max_new: int = 256, |
| temperature: float = 0.8, |
| top_p: float = 0.95, |
| top_k: int = 50, |
| ) -> Generator[Tuple[str, str], None, None]: |
| |
| if not prompt or not prompt.strip(): |
| yield "", "⚠️ Please enter a prompt" |
| return |
|
|
| domain_token = DOMAIN_TOKENS.get(domain, "<|domain_code|>") |
| full_prompt = domain_token + prompt.strip() |
|
|
| enc = tokenizer.encode(full_prompt) |
| ids = enc.ids |
| |
| if len(ids) > cfg.max_len - max_new: |
| ids = ids[-(cfg.max_len - max_new):] |
|
|
| x = torch.tensor([ids], dtype=torch.long) |
| generated = [] |
| t0 = time.time() |
|
|
| for step in range(int(max_new)): |
| x_in = x[:, -cfg.max_len:] if x.size(1) > cfg.max_len else x |
| logits = model(x_in)[0, -1, :] |
|
|
| logits = logits / max(float(temperature), 1e-5) |
|
|
| if top_k > 0: |
| k = min(int(top_k), logits.size(-1)) |
| topk_vals,_ = torch.topk(logits, k) |
| logits[logits < topk_vals[-1]] = float('-inf') |
|
|
| probs = F.softmax(logits, dim=-1) |
| sorted_p, sorted_idx = probs.sort(descending=True) |
| cumsum_p = sorted_p.cumsum(0) |
| sorted_p[(cumsum_p - sorted_p) > float(top_p)] = 0.0 |
| sorted_p = sorted_p / sorted_p.sum().clamp(min=1e-9) |
|
|
| next_tok = sorted_idx[torch.multinomial(sorted_p, num_samples=1)].item() |
|
|
| generated.append(next_tok) |
| x = torch.cat([x, torch.tensor([[next_tok]])], dim=1) |
|
|
| full_text = tokenizer.decode(ids + generated, skip_special_tokens=False) |
| |
| display_text = full_text |
| for tok in DOMAIN_TOKENS.values(): |
| display_text = display_text.replace(tok, "") |
|
|
| elapsed = time.time() - t0 |
| tps = len(generated) / max(elapsed, 1e-3) |
| stats = ( |
| f"⏱ {elapsed:.1f}s | " |
| f"🔤 {len(generated)}/{max_new} | " |
| f"⚡ {tps:.1f} tok/s | " |
| f"🎯 {domain}" |
| ) |
|
|
| yield display_text, stats |
|
|
| if EOS_ID >= 0 and next_tok == EOS_ID: |
| break |
|
|
| full_text = tokenizer.decode(ids + generated, skip_special_tokens=False) |
| for tok in DOMAIN_TOKENS.values(): |
| full_text = full_text.replace(tok, "") |
| |
| elapsed = time.time() - t0 |
| tps = len(generated) / max(elapsed, 1e-3) |
| yield ( |
| full_text, |
| f"✅ Done | {elapsed:.1f}s | {len(generated)} tokens | {tps:.1f} tok/s" |
| ) |
|
|
| |
| |
| |
| EXAMPLES = [ |
| ["def fibonacci(n):\n ", "Code", 64, 0.2, 0.95, 50], |
| ["1 + 1 = ", "Math", 10, 0.1, 0.95, 50], |
| ["class Stack:\n def __init__(self):\n ", "Code", 100, 0.3, 0.95, 50], |
| ["The Roman Empire began ", "General", 80, 0.6, 0.95, 50], |
| ] |
|
|
| with gr.Blocks(title="CodVa-2") as demo: |
| |
| gr.Markdown(""" |
| # 🧠 CodVa-2 — Pretrained Code LM |
| **213M params** · **2 vCPUs (free tier)** · Expect ~5-8 tok/s |
| |
| Pretrained model — continues text. Set **low max tokens** for short answers. |
| """) |
|
|
| with gr.Row(): |
| with gr.Column(scale=1): |
| prompt_box = gr.Textbox( |
| label="Prompt", |
| placeholder="def fibonacci(n):\n ", |
| lines=8, |
| ) |
| |
| domain_dropdown = gr.Dropdown( |
| choices=list(DOMAIN_TOKENS.keys()), |
| value="Code", |
| label="Domain", |
| ) |
| |
| with gr.Row(): |
| max_new_slider = gr.Slider(8, 256, value=64, step=8, label="Max tokens") |
| temp_slider = gr.Slider(0.0, 1.5, value=0.3, step=0.05, label="Temp") |
| |
| with gr.Row(): |
| gen_btn = gr.Button("▶ Generate", variant="primary", scale=2) |
| stop_btn = gr.Button("⏹ Stop", variant="stop", scale=1) |
|
|
| with gr.Column(scale=1): |
| output_box = gr.Textbox(label="Output", lines=15, interactive=False) |
| stats_box = gr.Textbox(label="", lines=1, interactive=False) |
|
|
| gr.Examples( |
| examples=EXAMPLES, |
| inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider], |
| label="Examples", |
| ) |
|
|
| gr.Markdown(""" |
| 💡 **Free tier = slow but it works!** |
| For `1+1=` type prompts: set max tokens to 10-15 and click Stop when done. |
| """) |
|
|
| gen_event = gen_btn.click( |
| fn=generate_stream, |
| inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider, |
| gr.State(0.95), gr.State(50)], |
| outputs=[output_box, stats_box], |
| ) |
| |
| stop_btn.click(fn=None, cancels=[gen_event]) |
|
|
| if __name__ == "__main__": |
| demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True) |
