########################################################################################################
# RWKV-7 "Goose" (x070 / g1d) HuggingFace configuration
# Based on the reference implementation from https://github.com/BlinkDL/RWKV-LM
########################################################################################################

from transformers.configuration_utils import PretrainedConfig


class RWKV7Config(PretrainedConfig):
    """Configuration for the RWKV-7 (x070 / g1d) language model.

    The defaults match the ``rwkv7-g1d-0.1b`` checkpoint
    (L12-D768, head_size 64), but the embedding / lm-head vocabulary has been
    re-sized (and re-initialized) to match the OLMo tokenizer.
    """

    model_type = "rwkv7"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size=100278,
        hidden_size=768,
        num_hidden_layers=12,
        head_size=64,
        intermediate_size=3072,
        decay_lora=64,
        aaa_lora=64,
        mv_lora=32,
        gate_lora=128,
        layer_norm_epsilon=1e-5,
        group_norm_epsilon=64e-5,
        bos_token_id=None,
        eos_token_id=100257,
        pad_token_id=100277,
        tie_word_embeddings=False,
        use_cuda_kernel=True,
        chunk_len=16,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.head_size = head_size
        self.intermediate_size = intermediate_size
        self.decay_lora = decay_lora
        self.aaa_lora = aaa_lora
        self.mv_lora = mv_lora
        self.gate_lora = gate_lora
        self.layer_norm_epsilon = layer_norm_epsilon
        self.group_norm_epsilon = group_norm_epsilon
        # attention/ffn dims are derived from hidden_size in the reference model
        self.attention_hidden_size = hidden_size
        self.use_cuda_kernel = use_cuda_kernel
        self.chunk_len = chunk_len

        assert hidden_size % head_size == 0, "hidden_size must be divisible by head_size"

        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )