"""Hugging Face configuration shared by every Weevil checkpoint."""

from transformers import PretrainedConfig


class WeevilConfig(PretrainedConfig):
    model_type = "weevil"

    def __init__(
        self, vocab_size=4096, hidden_size=256, embedding_dim=128,
        intermediate_size=960, num_hidden_layers=8, num_heads=8,
        qkv_rank=None, conv_kernel=4, norm_eps=1e-6, use_cache=True,
        **kwargs,
    ):
        kwargs.setdefault("bos_token_id", 0)
        kwargs.setdefault("eos_token_id", 0)
        kwargs.setdefault("pad_token_id", 1)
        kwargs.setdefault("tie_word_embeddings", False)
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.embedding_dim = embedding_dim
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_heads = num_heads
        self.qkv_rank = qkv_rank
        self.conv_kernel = conv_kernel
        self.norm_eps = norm_eps
        self.use_cache = use_cache
        if min(vocab_size, hidden_size, embedding_dim, intermediate_size,
               num_hidden_layers, num_heads, conv_kernel) < 1 or hidden_size < num_heads:
            raise ValueError("Weevil dimensions must be positive; heads cannot exceed hidden size")
        if qkv_rank is not None and not 1 <= qkv_rank <= hidden_size:
            raise ValueError("qkv_rank must be between 1 and hidden_size")


WeevilConfig.register_for_auto_class()