| |
|
|
| import os |
| import math |
| import time |
| import json |
| import threading |
|
|
| 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 |
| from dataclasses import dataclass |
| from tokenizers import Tokenizer |
| from huggingface_hub import hf_hub_download, login, HfApi |
| from flask import Flask, request, Response, send_file, jsonify |
| import re |
|
|
| torch.set_num_threads(2) |
| torch.set_num_interop_threads(1) |
|
|
| HF_TOKEN = os.environ.get("HF_TOKEN", "") |
| HF_DATA_REPO = "Bc-AI/MiniCoder-Instruction-dataset" |
| HF_MODEL_REPO = os.environ.get("MODEL_REPO", "hugging-science/CodVa-2-SFT") |
|
|
| 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, eps=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_DATA_REPO, |
| filename="codva2_sft_tokenizer.json", |
| repo_type="dataset", |
| token=HF_TOKEN or None, |
| ) |
| tokenizer = Tokenizer.from_file(tok_path) |
|
|
| IM_START_ID = tokenizer.token_to_id("<|im_start|>") |
| IM_END_ID = tokenizer.token_to_id("<|im_end|>") |
| EOS_ID = tokenizer.token_to_id("<|endoftext|>") or 50256 |
| STOP_IDS = {IM_END_ID, EOS_ID} |
|
|
| MAX_CTX = 2048 |
|
|
| 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)) |
| ckpts = sorted( |
| [f for f in files if "codva2_sft" in f and f.endswith(".pt")], |
| key=lambda x: int(x.split("step")[-1].split(".")[0]) if "step" in x else 0 |
| ) |
|
|
| if ckpts: |
| print(f"[init] loading: {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 SFT weights found") |
|
|
| model.eval() |
| print(f"[init] ready | {sum(p.numel() for p in model.parameters())/1e6:.1f}M params") |
|
|
| |
| |
| |
| gen_lock = threading.Lock() |
| stop_signal = threading.Event() |
|
|
| SYSTEM_PROMPT = "You are CodVa-2, a coding-specialized AI assistant. You write clean, correct, well-commented code. When solving problems, think step by step." |
|
|
| def format_prompt(user_msg: str) -> str: |
| return ( |
| f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n" |
| f"<|im_start|>user\n{user_msg}<|im_end|>\n" |
| f"<|im_start|>assistant\n" |
| ) |
|
|
| |
| |
| |
| app = Flask(__name__) |
|
|
| @app.route("/") |
| def index(): |
| return send_file("index.html") |
|
|
| @app.route("/stop", methods=["POST"]) |
| def stop(): |
| stop_signal.set() |
| return jsonify({"ok": True}) |
|
|
| @app.route("/generate", methods=["POST"]) |
| def generate(): |
| data = request.json |
| prompt = (data.get("prompt") or "").strip() |
| max_new = int(data.get("max_new", 256)) |
| temperature = float(data.get("temperature", 0.3)) |
| top_p = float(data.get("top_p", 0.95)) |
| top_k = int(data.get("top_k", 50)) |
|
|
| if not prompt: |
| return jsonify({"error": "empty prompt"}), 400 |
|
|
| stop_signal.clear() |
|
|
| def stream(): |
| full_prompt = format_prompt(prompt) |
| enc = tokenizer.encode(full_prompt) |
| ids = enc.ids |
| if len(ids) > MAX_CTX - max_new: |
| ids = ids[-(MAX_CTX - max_new):] |
|
|
| prompt_len = len(ids) |
| x = torch.tensor([ids], dtype=torch.long) |
| generated = [] |
| t0 = time.time() |
|
|
| |
| yield f"data: {json.dumps({'type': 'ctx', 'prompt_tokens': prompt_len, 'max_ctx': MAX_CTX, 'max_new': max_new})}\n\n" |
|
|
| with torch.no_grad(): |
| for step in range(max_new): |
| if stop_signal.is_set(): |
| break |
|
|
| x_in = x[:, -MAX_CTX:] if x.size(1) > MAX_CTX else x |
| logits = model(x_in)[0, -1, :] |
|
|
| |
| logits = logits / max(temperature, 1e-5) |
| if top_k > 0: |
| k = min(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 = sorted_p.cumsum(0) |
| sorted_p[(cumsum - sorted_p) > top_p] = 0.0 |
| sorted_p = sorted_p / sorted_p.sum().clamp(min=1e-9) |
| next_tok = sorted_idx[torch.multinomial(sorted_p, 1)].item() |
|
|
| generated.append(next_tok) |
| x = torch.cat([x, torch.tensor([[next_tok]])], dim=1) |
|
|
| |
| raw = tokenizer.decode(generated, skip_special_tokens=False) |
|
|
| elapsed = time.time() - t0 |
| tps = len(generated) / max(elapsed, 1e-3) |
| total_tokens = prompt_len + len(generated) |
|
|
| payload = { |
| "type": "token", |
| "raw": raw, |
| "tok_id": next_tok, |
| "step": len(generated), |
| "max_new": max_new, |
| "tps": round(tps, 1), |
| "elapsed": round(elapsed, 1), |
| "total_tokens": total_tokens, |
| "prompt_tokens": prompt_len, |
| "max_ctx": MAX_CTX, |
| } |
| yield f"data: {json.dumps(payload)}\n\n" |
|
|
| if next_tok in STOP_IDS: |
| break |
|
|
| elapsed = time.time() - t0 |
| tps = len(generated) / max(elapsed, 1e-3) |
| yield f"data: {json.dumps({'type': 'done', 'tps': round(tps,1), 'elapsed': round(elapsed,1), 'tokens': len(generated)})}\n\n" |
|
|
| return Response(stream(), mimetype="text/event-stream", |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) |
|
|
| if __name__ == "__main__": |
| app.run(host="0.0.0.0", port=7860, threaded=True) |
