"""
Configuration for LunarisGuardModel.

This config is intentionally minimal — it just declares the dual-head
classification setup on top of ModernBERT-base. The backbone config
itself comes from the underlying AutoConfig when the model is loaded.
"""

from transformers import PretrainedConfig


class LunarisGuardConfig(PretrainedConfig):
    """Config for the dual-head Lunaris Guard classifier."""

    model_type = "lunaris_guard"

    def __init__(
        self,
        base_model_name: str = "answerdotai/ModernBERT-base",
        hidden_size: int = 768,
        num_injection_classes: int = 2,
        num_safety_classes: int = 2,
        classifier_dropout: float = 0.15,
        max_length: int = 2048,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.base_model_name = base_model_name
        self.hidden_size = hidden_size
        self.num_injection_classes = num_injection_classes
        self.num_safety_classes = num_safety_classes
        self.classifier_dropout = classifier_dropout
        self.max_length = max_length

        # Standard HF id2label / label2id for Hub display
        self.id2label_injection = {0: "BENIGN", 1: "INJECTION"}
        self.label2id_injection = {"BENIGN": 0, "INJECTION": 1}
        self.id2label_safety = {0: "SAFE", 1: "UNSAFE"}
        self.label2id_safety = {"SAFE": 0, "UNSAFE": 1}