modeling_codva1.py

"""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!
""")