"""Hugging Face configuration for compliantLLM."""

from transformers import PretrainedConfig


class CompliantLLMConfig(PretrainedConfig):
    model_type = "compliant_llm"

    def __init__(
        self,
        input_vocab_size=256,
        output_vocab_size=3,
        max_context=1024,
        output_length=3,
        d_model=64,
        n_heads=4,
        n_layers=2,
        ffn_dim=128,
        dropout=0.0,
        output_tokens=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.input_vocab_size = input_vocab_size
        self.output_vocab_size = output_vocab_size
        self.max_context = max_context
        self.output_length = output_length
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.ffn_dim = ffn_dim
        self.dropout = dropout
        self.output_tokens = output_tokens or [
            "Sorry, but that question violates GDPR.",
            "<|end_turn|>",
            "<|eos|>",
        ]

        if self.input_vocab_size != 256:
            raise ValueError("compliantLLM requires exactly 256 input tokens")
        if self.output_vocab_size != 3 or self.output_length != 3:
            raise ValueError("compliantLLM requires exactly three output tokens and positions")
        if self.max_context != 1024:
            raise ValueError("compliantLLM requires a 1024-token context")