"""CodVa-1 Model for HuggingFace — Dense-CodeMoE causal LM. Architecture (faithful port of the original training script's `CodvaModel`): - RMSNorm pre-norm + residual - GQA attention with QK-norm + RoPE - SwiGLU dense FFN on non-MoE layers - CodeMoE FFN (token-level top-k routing + always-on shared experts + aux-loss-free load balancing) on MoE layers, every `moe_every` layers - Optional ablatable AST-structural attention bias - Weight tying NOTE: This model does not implement a KV cache. `use_cache` is always forced to False, matching the Nova-1 HF wrapper convention. """ from __future__ import annotations from typing import Optional import warnings import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast try: from .configuration_codva1 import CodVa1Config except ImportError: from configuration_codva1 import CodVa1Config # ══════════════════════════════════════════════════════════════════════════════ # HELPERS # ══════════════════════════════════════════════════════════════════════════════ def _precompute_rope(head_dim, max_len, theta=10000.0, device=None, dtype=torch.float32): inv_freq = 1.0 / (theta ** ( torch.arange(0, head_dim, 2, device=device, dtype=dtype) / head_dim )) t = torch.arange(max_len, device=device, dtype=dtype) freqs = torch.outer(t, inv_freq) return freqs.cos(), freqs.sin() def _apply_rope(xq, xk, cos, sin): """ xq, xk: [B, L, n_heads, head_dim] (pre-transpose layout, matching the original training script's GQAAttention, which applies RoPE BEFORE transposing heads to the front). """ L = xq.shape[1] cos = cos.to(device=xq.device)[:L] sin = sin.to(device=xq.device)[:L] c = torch.cat([cos, cos], dim=-1).unsqueeze(0).unsqueeze(2) # [1, L, 1, hd] s = torch.cat([sin, sin], dim=-1).unsqueeze(0).unsqueeze(2) def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat([-x2, x1], dim=-1) xq_f, xk_f = xq.float(), xk.float() xq_out = (xq_f * c + rotate_half(xq_f) * s).to(xq.dtype) xk_out = (xk_f * c + rotate_half(xk_f) * s).to(xk.dtype) return xq_out, xk_out # ══════════════════════════════════════════════════════════════════════════════ # MODULES # ══════════════════════════════════════════════════════════════════════════════ class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.scale = nn.Parameter(torch.ones(dim)) def forward(self, x): return F.rms_norm(x, self.scale.shape, self.scale, self.eps) class SwiGLU(nn.Module): def __init__(self, d_model: int, hidden: int): super().__init__() self.gate = nn.Linear(d_model, hidden, bias=False) self.up = nn.Linear(d_model, hidden, bias=False) self.down = nn.Linear(hidden, d_model, bias=False) def forward(self, x): return self.down(F.silu(self.gate(x)) * self.up(x)) class GQAAttention(nn.Module): def __init__(self, config: CodVa1Config): super().__init__() assert config.d_model % config.n_heads == 0 assert config.n_heads % config.n_kv_heads == 0 self.nh = config.n_heads self.nkv = config.n_kv_heads self.hd = config.d_model // config.n_heads self.ng = config.n_heads // config.n_kv_heads self.use_qk_norm = config.use_qk_norm D, Dkv = config.n_heads * self.hd, config.n_kv_heads * self.hd self.q = nn.Linear(config.d_model, D, bias=False) self.k = nn.Linear(config.d_model, Dkv, bias=False) self.v = nn.Linear(config.d_model, Dkv, bias=False) self.o = nn.Linear(D, config.d_model, bias=False) if self.use_qk_norm: self.q_norm = RMSNorm(self.hd, config.rms_norm_eps) self.k_norm = RMSNorm(self.hd, config.rms_norm_eps) def forward(self, x, cos, sin, attn_bias=None): B, L, _ = x.shape q = self.q(x).view(B, L, self.nh, self.hd) k = self.k(x).view(B, L, self.nkv, self.hd) v = self.v(x).view(B, L, self.nkv, self.hd) if self.use_qk_norm: q = self.q_norm(q) k = self.k_norm(k) q, k = _apply_rope(q, k, cos, sin) q = q.transpose(1, 2) # [B, nh, L, hd] k = k.transpose(1, 2).repeat_interleave(self.ng, dim=1) # [B, nh, L, hd] v = v.transpose(1, 2).repeat_interleave(self.ng, dim=1) if attn_bias is not None: causal = torch.triu( torch.ones(L, L, device=x.device, dtype=torch.bool), diagonal=1) bias = attn_bias.masked_fill(causal.view(1, 1, L, L), float('-inf')) out = F.scaled_dot_product_attention( q, k, v, attn_mask=bias, is_causal=False, dropout_p=0.0) else: out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0) return self.o(out.transpose(1, 2).contiguous().view(B, L, -1)) class StructuralBias(nn.Module): """Optional AST-structural relation bias added to attention logits.""" def __init__(self, n_heads: int, n_rel: int): super().__init__() self.n_rel = n_rel + 1 # 0 == "no relation" self.rel_emb = nn.Parameter(torch.zeros(n_heads, self.n_rel)) def forward(self, rel_ids: torch.Tensor) -> torch.Tensor: # rel_ids: [B, L, L] long -> [B, nh, L, L] bias = self.rel_emb.t()[rel_ids] return bias.permute(0, 3, 1, 2).contiguous() class CodeMoEFFN(nn.Module): """ Token-level sparse MoE FFN (Mixtral/DeepSeek-MoE style): - shared expert(s) always active - top-k routed experts, dispatched via sort/permute (no Python loop over tokens — only over the (small) number of experts) - aux-loss-free load balancing bias (DeepSeek-V3 trick). The bias only drifts during `.train()`; at inference (`.eval()`) it is static and simply added to the router logits. """ def __init__(self, config: CodVa1Config): super().__init__() self.d_model = config.d_model self.top_k = config.moe_top_k self.n_exp = config.moe_experts self.bias_speed = config.moe_bias_speed self.shared = nn.ModuleList( [SwiGLU(config.d_model, config.moe_hidden) for _ in range(config.moe_shared)]) self.experts = nn.ModuleList( [SwiGLU(config.d_model, config.moe_hidden) for _ in range(config.moe_experts)]) self.gate = nn.Linear(config.d_model, config.moe_experts, bias=False) self.register_buffer("balance_bias", torch.zeros(config.moe_experts)) self.register_buffer("load_ema", torch.zeros(config.moe_experts)) def forward(self, x: torch.Tensor) -> torch.Tensor: B, L, D = x.shape xf = x.reshape(-1, D) N = xf.shape[0] logits = self.gate(xf) biased = logits + self.balance_bias _, idx = torch.topk(biased, self.top_k, dim=-1) w = logits.softmax(-1).gather(1, idx) w = w / (w.sum(-1, keepdim=True) + 1e-9) out = torch.zeros_like(xf) for s in self.shared: out = out + s(xf) flat_idx = idx.reshape(-1) flat_w = w.reshape(-1) token_ids = torch.arange(N, device=x.device)\ .unsqueeze(1).expand(-1, self.top_k).reshape(-1) sort_order = flat_idx.argsort(stable=True) sorted_exp = flat_idx[sort_order] sorted_tok = token_ids[sort_order] sorted_w = flat_w[sort_order] counts = torch.bincount(sorted_exp, minlength=self.n_exp) dispatched = xf[sorted_tok] routed_out = torch.zeros_like(dispatched) start = 0 for e in range(self.n_exp): end = start + int(counts[e]) if end > start: routed_out[start:end] = self.experts[e](dispatched[start:end]) start = end routed_out = routed_out * sorted_w.unsqueeze(-1) unsort_order = sort_order.argsort() routed_out = routed_out[unsort_order] routed_out = routed_out.view(N, self.top_k, D) out = out + routed_out.sum(dim=1) if self.training: with torch.no_grad(): frac = counts.float() / (counts.sum() + 1e-9) target = 1.0 / self.n_exp self.balance_bias.add_(self.bias_speed * (target - frac)) self.load_ema.mul_(0.99).add_(frac, alpha=0.01) return out.view(B, L, D) class CodVa1Block(nn.Module): def __init__(self, config: CodVa1Config, layer_idx: int): super().__init__() self.attn_norm = RMSNorm(config.d_model, config.rms_norm_eps) self.ffn_norm = RMSNorm(config.d_model, config.rms_norm_eps) self.attn = GQAAttention(config) self.is_moe = config.use_moe and (layer_idx % config.moe_every == 0) self.ffn = (CodeMoEFFN(config) if self.is_moe else SwiGLU(config.d_model, config.ffn_hidden)) def forward(self, x, cos, sin, attn_bias=None): x = x + self.attn(self.attn_norm(x), cos, sin, attn_bias=attn_bias) x = x + self.ffn(self.ffn_norm(x)) return x # ══════════════════════════════════════════════════════════════════════════════ # MAIN MODEL # ══════════════════════════════════════════════════════════════════════════════ class CodVa1ForCausalLM(PreTrainedModel, GenerationMixin): config_class = CodVa1Config supports_gradient_checkpointing = True _supports_cache_class = False def __init__(self, config: CodVa1Config): super().__init__(config) self.embed = nn.Embedding(config.vocab_size, config.d_model) self.layers = nn.ModuleList( [CodVa1Block(config, i) for i in range(config.n_layers)]) self.norm = RMSNorm(config.d_model, config.rms_norm_eps) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) self.struct_bias = (StructuralBias(config.n_heads, config.n_struct_rel) if config.use_structural_bias else None) self._rope_cache: Optional[tuple] = None self.post_init() def get_input_embeddings(self): return self.embed def set_input_embeddings(self, v): self.embed = v def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, v): self.lm_head = v def tie_weights(self, *args, **kwargs): if getattr(self.config, "tie_word_embeddings", True): self.lm_head.weight = self.embed.weight try: super().tie_weights(*args, **kwargs) except TypeError: super().tie_weights() def _get_rope(self, T: int, device, dtype): cache = self._rope_cache if cache is not None: cT, cdev, cdt, ccos, csin = cache if cT >= T and cdev == device and cdt == dtype: return ccos[:T], csin[:T] head_dim = self.config.d_model // self.config.n_heads alloc = max(T, self.config.max_len) cos, sin = _precompute_rope( head_dim, alloc, self.config.rope_theta, device=device, dtype=dtype ) self._rope_cache = (alloc, device, dtype, cos, sin) return cos[:T], sin[:T] def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[list] = None, rel_ids: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ) -> CausalLMOutputWithPast: # ─── NO KV CACHE SUPPORT ─── if use_cache: warnings.warn( "`CodVa1ForCausalLM` does not support KV cache. Forcing `use_cache=False`.", UserWarning ) if past_key_values is not None: warnings.warn( "`past_key_values` were passed but `CodVa1ForCausalLM` does not support KV cache. Ignoring.", UserWarning ) past_key_values = None # ──────────────────────────── B, T = input_ids.shape x = self.embed(input_ids) cos, sin = self._get_rope(T, x.device, torch.float32) attn_bias = None if self.struct_bias is not None and rel_ids is not None: sb_device = next(self.struct_bias.parameters()).device attn_bias = self.struct_bias(rel_ids.to(sb_device)) for layer in self.layers: layer_device = next(layer.parameters()).device x = x.to(layer_device) cos_l = cos.to(layer_device) sin_l = sin.to(layer_device) bias_l = attn_bias.to(layer_device) if attn_bias is not None else None x = layer(x, cos_l, sin_l, attn_bias=bias_l) norm_device = next(self.norm.parameters()).device x = self.norm(x.to(norm_device)) logits = self.lm_head(x) loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous().to(logits.device) loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)).float(), shift_labels.view(-1), ignore_index=-100, ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=None, ) def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs): # No KV cache: always feed the full sequence. if past_key_values is not None: past_key_values = None if input_ids.size(1) > self.config.max_len: input_ids = input_ids[:, -self.config.max_len:] return {"input_ids": input_ids, "attention_mask": attention_mask} print(""" 👋 Hey! Quick message from the SmilyAI-Labs team! CodVa was designed specifically for coding tasks — and that's where it absolutely shines 🔥 Here's how to get the best out of it: 🏷️ DOMAIN TAGS — slap one of these at the start of your prompt: <|domain_code|> → general coding tasks (default for most stuff) <|domain_math|> → math / algorithms <|domain_general|> → general text / comments / docs <|domain_reasoning|> → step-by-step reasoning / problem solving 💻 CODE TAGS — wrap your code prompts like this: <|domain_code|><|code_start|>def your_function(): 📐 MATH TAGS — for math problems: <|domain_math|><|math_start|>solve for x... 💬 CHAT FORMAT — for instruction style prompts: <|im_start|>user your message here <|im_end|> <|im_start|>assistant 🚀 QUICK START EXAMPLE: from transformers import pipeline pipe = pipeline("text-generation", model="Bc-AI/codva-checkpoints", trust_remote_code=True) pipe("<|domain_code|><|code_start|>def fibonacci(n):", max_new_tokens=200) CodVa was pretrained on a massive code corpus and fine-tuned to understand code structure deeply — so the more context you give it the better it performs! Give it a full function signature, a docstring, some examples — it will fill in the rest like a champ 💪 Happy coding! — SmilyAI-Labs 🧪 P.S, CodVa is best at coding, for general tasks, try our Nova model! """)