model/model.py

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)