"""Hugging Face configuration for Needle."""

from __future__ import annotations

from transformers import PretrainedConfig


class NeedleConfig(PretrainedConfig):
    model_type = "needle"

    def __init__(
        self,
        vocab_size: int = 8192,
        hidden_size: int | None = None,
        d_model: int = 512,
        num_attention_heads: int | None = None,
        num_heads: int = 8,
        num_key_value_heads: int | None = None,
        num_kv_heads: int = 4,
        num_encoder_layers: int = 12,
        num_decoder_layers: int = 8,
        rope_theta: float = 10000.0,
        rms_norm_eps: float = 1e-6,
        pad_token_id: int = 0,
        eos_token_id: int = 1,
        bos_token_id: int = 2,
        unk_token_id: int = 3,
        decoder_start_token_id: int | None = None,
        tie_word_embeddings: bool = True,
        torch_dtype: str = "bfloat16",
        **kwargs,
    ) -> None:
        kwargs.pop("is_encoder_decoder", None)
        hidden_size = int(hidden_size if hidden_size is not None else d_model)
        num_attention_heads = int(num_attention_heads if num_attention_heads is not None else num_heads)
        num_key_value_heads = int(num_key_value_heads if num_key_value_heads is not None else num_kv_heads)
        decoder_start_token_id = eos_token_id if decoder_start_token_id is None else decoder_start_token_id

        super().__init__(
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
            bos_token_id=bos_token_id,
            unk_token_id=unk_token_id,
            decoder_start_token_id=decoder_start_token_id,
            tie_word_embeddings=tie_word_embeddings,
            is_encoder_decoder=True,
            torch_dtype=torch_dtype,
            **kwargs,
        )

        self.vocab_size = int(vocab_size)
        self.hidden_size = hidden_size
        self.d_model = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.num_kv_heads = num_key_value_heads
        self.num_encoder_layers = int(num_encoder_layers)
        self.num_decoder_layers = int(num_decoder_layers)
        self.num_hidden_layers = int(num_decoder_layers)
        self.rope_theta = float(rope_theta)
        self.rms_norm_eps = float(rms_norm_eps)
        self.attention_head_dim = hidden_size // max(1, num_attention_heads)