import math from dataclasses import asdict from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from .config import GuardModelConfig GUARD_LABELS = { 0: "SAFE", 1: "JAILBREAK_ATTACK", 2: "PROMPT_INJECTION", 3: "TOXIC_HARASSMENT", 4: "PII_SENSITIVE_LEAK", 5: "MALICIOUS_INTENT", } class RMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: norm = x.pow(2).mean(dim=-1, keepdim=True) return self.weight * x * torch.rsqrt(norm + self.eps) def rotate_half(x: torch.Tensor) -> torch.Tensor: x1 = x[..., ::2] x2 = x[..., 1::2] return torch.stack((-x2, x1), dim=-1).flatten(-2) class RotaryEmbedding(nn.Module): def __init__(self, head_dim: int, base: float = 500000.0) -> None: super().__init__() inv_freq = base ** (-torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim) self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, q: torch.Tensor, k: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: angles = position_ids[:, None, :, None] * self.inv_freq[None, None, None, :] cos = torch.repeat_interleave(torch.cos(angles), 2, dim=-1).to(dtype=q.dtype) sin = torch.repeat_interleave(torch.sin(angles), 2, dim=-1).to(dtype=q.dtype) return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin class GroupedQueryAttention(nn.Module): def __init__(self, cfg: GuardModelConfig) -> None: super().__init__() self.num_heads = cfg.num_heads self.num_kv_heads = cfg.num_kv_heads self.head_dim = cfg.head_dim self.hidden_size = cfg.hidden_size self.num_queries_per_kv = self.num_heads // self.num_kv_heads self.q_proj = nn.Linear(cfg.hidden_size, cfg.num_heads * cfg.head_dim, bias=False) self.k_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * cfg.head_dim, bias=False) self.v_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * cfg.head_dim, bias=False) self.o_proj = nn.Linear(cfg.num_heads * cfg.head_dim, cfg.hidden_size, bias=False) self.rotary = RotaryEmbedding(cfg.head_dim, base=cfg.rope_base) def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: bsz, seq_len, _ = x.shape q = self.q_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) q, k = self.rotary(q, k, position_ids=position_ids) if self.num_queries_per_kv > 1: k = k.unsqueeze(2).expand(bsz, self.num_kv_heads, self.num_queries_per_kv, seq_len, self.head_dim).reshape(bsz, self.num_heads, seq_len, self.head_dim) v = v.unsqueeze(2).expand(bsz, self.num_kv_heads, self.num_queries_per_kv, seq_len, self.head_dim).reshape(bsz, self.num_heads, seq_len, self.head_dim) out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=True) out = out.transpose(1, 2).contiguous().view(bsz, seq_len, self.hidden_size) return self.o_proj(out) class SwiGLU(nn.Module): def __init__(self, hidden_size: int, intermediate_size: int) -> None: super().__init__() self.gate = nn.Linear(hidden_size, intermediate_size, bias=False) self.up = nn.Linear(hidden_size, intermediate_size, bias=False) self.down = nn.Linear(intermediate_size, hidden_size, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down(F.silu(self.gate(x)) * self.up(x)) class GuardBlock(nn.Module): def __init__(self, cfg: GuardModelConfig) -> None: super().__init__() self.norm1 = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) self.attn = GroupedQueryAttention(cfg) self.norm2 = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) self.ffn = SwiGLU(cfg.hidden_size, cfg.intermediate_size) def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.norm1(x), position_ids=position_ids) x = x + self.ffn(self.norm2(x)) return x class CortiqaGuard40M(nn.Module): """ Cortiqa Guard 40M: Ultra-fast On-Device LLM Firewall & Safety Guardrail. Dual Architecture: 1. Generative LM Head: Explains why prompt was blocked or generates tokens. 2. Fast Classification Head: Single-pass logits output for <3ms firewall latency. """ def __init__(self, cfg: GuardModelConfig) -> None: super().__init__() self.cfg = cfg self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size) self.layers = nn.ModuleList([GuardBlock(cfg) for _ in range(cfg.num_layers)]) self.final_norm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) # 1. Generative LM Head (Tied) self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) if cfg.tie_embeddings: self.lm_head.weight = self.embed_tokens.weight # 2. Fast Classifier Head (0: SAFE, 1: JAILBREAK, 2: INJECTION, 3: TOXIC, 4: PII, 5: MALICIOUS) self.classifier_head = nn.Sequential( nn.Linear(cfg.hidden_size, cfg.hidden_size), nn.SiLU(), nn.Linear(cfg.hidden_size, cfg.num_classes), ) self.apply(self._init_weights) def _init_weights(self, module: nn.Module) -> None: if isinstance(module, (nn.Linear, nn.Embedding)): nn.init.normal_(module.weight, mean=0.0, std=0.02) if isinstance(module, nn.Linear) and module.bias is not None: nn.init.zeros_(module.bias) def forward( self, input_ids: torch.Tensor, labels: Optional[torch.Tensor] = None, class_labels: Optional[torch.Tensor] = None, ) -> Dict[str, torch.Tensor]: bsz, seq_len = input_ids.shape x = self.embed_tokens(input_ids) position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(bsz, -1) for layer in self.layers: x = layer(x, position_ids=position_ids) x = self.final_norm(x) # 1. Causal LM Logits lm_logits = self.lm_head(x) # 2. Classification Logits (Pooled from last active token) pooled_repr = x[:, -1, :] class_logits = self.classifier_head(pooled_repr) out = {"lm_logits": lm_logits, "class_logits": class_logits} total_loss = 0.0 if class_labels is not None: cls_loss = F.cross_entropy(class_logits, class_labels) out["cls_loss"] = cls_loss total_loss = total_loss + cls_loss if labels is not None: shift_logits = lm_logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() lm_loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100) out["lm_loss"] = lm_loss total_loss = total_loss + 0.5 * lm_loss if class_labels is not None or labels is not None: out["loss"] = total_loss return out @torch.no_grad() def inspect_prompt(self, input_ids: torch.Tensor) -> Dict[str, object]: """ Fast Real-time Firewall Interceptor (runs in <3ms on CPU/Edge). Returns classification label, confidence score, and is_safe boolean. """ self.eval() bsz, seq_len = input_ids.shape x = self.embed_tokens(input_ids) position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(bsz, -1) for layer in self.layers: x = layer(x, position_ids=position_ids) x = self.final_norm(x) class_logits = self.classifier_head(x[:, -1, :]) probs = F.softmax(class_logits, dim=-1)[0] pred_class = int(torch.argmax(probs).item()) confidence = float(probs[pred_class].item()) is_safe = (pred_class == 0) label_str = GUARD_LABELS.get(pred_class, "UNKNOWN") return { "is_safe": is_safe, "label": label_str, "class_id": pred_class, "confidence": confidence, "all_scores": {GUARD_LABELS[i]: float(probs[i].item()) for i in range(len(GUARD_LABELS))}, } def param_count(self) -> int: return sum(p.numel() for p in self.parameters() if p.requires_grad)