Upload modeling_lunaris_guard.py

modeling_lunaris_guard.py

@@ -0,0 +1,76 @@
+"""
+LunarisGuardModel — Hub-loadable version.
+
+This module is what gets loaded when a user runs:
+
+    AutoModel.from_pretrained("auren-research/lunaris-guard", trust_remote_code=True)
+
+The class is the same dual-head classifier used in training, but configured
+to load weights from a HF Hub repo via PretrainedConfig + auto_map.
+
+Returns a dict with keys:
+  - injection_logits: [B, 2]
+  - safety_logits:    [B, 2]
+  - pooled_output:    [B, hidden_size]   (debug / probing)
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from transformers import AutoModel, PreTrainedModel
+
+from .configuration_lunaris_guard import LunarisGuardConfig
+
+
+class LunarisGuardModel(PreTrainedModel):
+    """Dual-head classifier: injection + content safety, on a ModernBERT backbone."""
+
+    config_class = LunarisGuardConfig
+    base_model_prefix = "backbone"
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: LunarisGuardConfig):
+        super().__init__(config)
+        self.config = config
+
+        # Backbone: ModernBERT-base. Loaded fresh; the saved state_dict will
+        # overwrite these weights when from_pretrained() runs.
+        self.backbone = AutoModel.from_pretrained(
+            config.base_model_name, trust_remote_code=True,
+        )
+
+        self.dropout = nn.Dropout(config.classifier_dropout)
+
+        self.injection_head = nn.Linear(
+            config.hidden_size, config.num_injection_classes
+        )
+        self.safety_head = nn.Linear(
+            config.hidden_size, config.num_safety_classes
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        injection_labels: Optional[torch.Tensor] = None,
+        safety_labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        outputs = self.backbone(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            return_dict=True,
+        )
+        # CLS pooling
+        pooled = outputs.last_hidden_state[:, 0, :]
+        pooled = self.dropout(pooled)
+
+        injection_logits = self.injection_head(pooled)
+        safety_logits = self.safety_head(pooled)
+
+        return {
+            "injection_logits": injection_logits,
+            "safety_logits": safety_logits,
+            "pooled_output": pooled,
+        }