"""AQ modeling code — Zyora Labs' proprietary, from-scratch decoder architecture. AQ (Academic Quotient) is a concept-first academic language model built from scratch in pure PyTorch. Architecture: RMSNorm, rotary position embeddings (RoPE), grouped-query attention (GQA), SwiGLU feed-forward, tied embeddings. Module names intentionally mirror the original AQ training code, so trained checkpoints load 1:1 with no key remapping. """ from __future__ import annotations from typing import Optional 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 CausalLMOutput from .configuration_aq import AQConfig class AQRMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: dtype = x.dtype x = x.float() x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return (x.to(dtype)) * self.weight def build_rope_cache(seq_len: int, head_dim: int, theta: float, device): inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) t = torch.arange(seq_len, device=device).float() freqs = torch.outer(t, inv_freq) emb = torch.cat((freqs, freqs), dim=-1) return emb.cos(), emb.sin() def rotate_half(x: torch.Tensor) -> torch.Tensor: x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) def apply_rope(q, k, cos, sin): cos = cos.unsqueeze(0).unsqueeze(0) sin = sin.unsqueeze(0).unsqueeze(0) q = (q * cos) + (rotate_half(q) * sin) k = (k * cos) + (rotate_half(k) * sin) return q, k class AQAttention(nn.Module): def __init__(self, cfg: AQConfig): super().__init__() self.n_heads = cfg.num_attention_heads self.n_kv = cfg.num_key_value_heads self.head_dim = cfg.head_dim self.n_rep = self.n_heads // self.n_kv self.q_proj = nn.Linear(cfg.hidden_size, self.n_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(cfg.hidden_size, self.n_kv * self.head_dim, bias=False) self.v_proj = nn.Linear(cfg.hidden_size, self.n_kv * self.head_dim, bias=False) self.o_proj = nn.Linear(self.n_heads * self.head_dim, cfg.hidden_size, bias=False) def forward(self, x, cos, sin, attn_mask: Optional[torch.Tensor] = None): B, T, _ = x.shape q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, T, self.n_kv, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, T, self.n_kv, self.head_dim).transpose(1, 2) q, k = apply_rope(q, k, cos, sin) k = k.repeat_interleave(self.n_rep, dim=1) v = v.repeat_interleave(self.n_rep, dim=1) if attn_mask is not None: out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) else: out = F.scaled_dot_product_attention(q, k, v, is_causal=True) out = out.transpose(1, 2).contiguous().view(B, T, -1) return self.o_proj(out) class AQSwiGLU(nn.Module): def __init__(self, cfg: AQConfig): super().__init__() self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class AQBlock(nn.Module): def __init__(self, cfg: AQConfig): super().__init__() self.attn_norm = AQRMSNorm(cfg.hidden_size, cfg.rms_norm_eps) self.attn = AQAttention(cfg) self.mlp_norm = AQRMSNorm(cfg.hidden_size, cfg.rms_norm_eps) self.mlp = AQSwiGLU(cfg) def forward(self, x, cos, sin, attn_mask=None): x = x + self.attn(self.attn_norm(x), cos, sin, attn_mask) x = x + self.mlp(self.mlp_norm(x)) return x class AQPreTrainedModel(PreTrainedModel): config_class = AQConfig base_model_prefix = "aq" supports_gradient_checkpointing = False _no_split_modules = ["AQBlock"] def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) class AQForCausalLM(AQPreTrainedModel, GenerationMixin): """AQ decoder language model with a causal LM head (tied embeddings).""" _tied_weights_keys = ["lm_head.weight"] def __init__(self, config: AQConfig): super().__init__(config) self.embed = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList([AQBlock(config) for _ in range(config.num_hidden_layers)]) self.norm = AQRMSNorm(config.hidden_size, config.rms_norm_eps) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) cos, sin = build_rope_cache( config.max_position_embeddings, config.head_dim, config.rope_theta, "cpu") self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.post_init() # weight init + embedding tying (config.tie_word_embeddings) def get_input_embeddings(self): return self.embed def set_input_embeddings(self, value): self.embed = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, **kwargs, ) -> CausalLMOutput: B, T = input_ids.shape dtype = self.embed.weight.dtype cos = self.rope_cos[:T].to(device=input_ids.device, dtype=dtype) sin = self.rope_sin[:T].to(device=input_ids.device, dtype=dtype) # Combined causal + padding mask (only when padding is actually present). attn_mask = None if attention_mask is not None and not bool(attention_mask.all()): causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=input_ids.device)) pad = attention_mask[:, None, None, :].to(torch.bool) # (B,1,1,T) attn_mask = causal[None, None, :, :] & pad x = self.embed(input_ids) for layer in self.layers: x = layer(x, cos, sin, attn_mask) x = self.norm(x) logits = self.lm_head(x) loss = None if labels is not None: shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100, ) return CausalLMOutput(loss=loss, logits=logits)