"""Hugging Face configuration for TaoNet."""

from transformers import PretrainedConfig


class TaoNetConfig(PretrainedConfig):
    """Configuration for the TaoNet causal language model."""

    model_type = "taonet"

    def __init__(
        self,
        vocab_size=8192,
        hidden_dim=768,
        num_layers=12,
        num_heads=12,
        head_dim=None,
        intermediate_dim=None,
        dropout=0.1,
        max_seq_length=2048,
        d_latent_kv=None,
        d_rope=None,
        gqa_groups=1,
        hidden_dim_ff=None,
        use_factorized_embedding=False,
        d_embed_rank=96,
        rope_scale=40.0,
        yarn_enabled=False,
        yarn_original_max_seq_length=None,
        yarn_alpha=1.0,
        init_std=0.02,
        pad_token_id=3,
        bos_token_id=1,
        eos_token_id=2,
        unk_token_id=0,
        **kwargs,
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            unk_token_id=unk_token_id,
            **kwargs,
        )
        self.vocab_size = vocab_size
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim if head_dim is not None else hidden_dim // num_heads
        self.intermediate_dim = intermediate_dim if intermediate_dim is not None else hidden_dim * 4
        self.dropout = dropout
        self.max_seq_length = max_seq_length
        self.d_latent_kv = d_latent_kv
        self.d_rope = d_rope
        self.gqa_groups = gqa_groups
        self.hidden_dim_ff = hidden_dim_ff
        self.use_factorized_embedding = use_factorized_embedding
        self.d_embed_rank = d_embed_rank
        self.rope_scale = rope_scale
        self.yarn_enabled = yarn_enabled
        self.yarn_original_max_seq_length = yarn_original_max_seq_length
        self.yarn_alpha = yarn_alpha
        self.init_std = init_std

    @classmethod
    def from_taotrain_model_config(cls, model_config, **extra_kwargs):
        """Build an HF config from a TaoTrain model-config dictionary."""
        data = dict(model_config)
        data.pop("architecture_type", None)
        data.pop("llm_architecture_type", None)
        data.update(extra_kwargs)
        return cls(**data)