from transformers import PretrainedConfig


class BETConfig(PretrainedConfig):
    model_type = "bet"

    def __init__(
        self,
        vocab_size=259,
        hidden_size=324,
        intermediate_size=864,
        prelude_layers=1,
        body_blocks=6,
        coda_layers=1,
        num_attention_heads=6,
        num_key_value_heads=2,
        head_dim=54,
        lora_rank=16,
        hyper_lanes=2,
        max_position_embeddings=1024,
        max_loops=8,
        rope_theta=10_000.0,
        rms_norm_eps=1e-6,
        ddl_beta_init=1.0,
        ddl_k_eps=1e-2,
        ddl_v_sigmoid_scale=4.0,
        refinement_cycles=8,
        use_cache=False,
        tie_word_embeddings=True,
        pad_token_id=256,
        bos_token_id=257,
        eos_token_id=258,
        **kwargs,
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            is_encoder_decoder=False,
            **kwargs,
        )
        self.vocab_size=int(vocab_size)
        self.hidden_size=int(hidden_size)
        self.intermediate_size=int(intermediate_size)
        self.prelude_layers=int(prelude_layers)
        self.body_blocks=int(body_blocks)
        self.coda_layers=int(coda_layers)
        # Common HF tooling expects num_hidden_layers even though only the body loops.
        self.num_hidden_layers=int(prelude_layers+body_blocks+coda_layers)
        self.num_attention_heads=int(num_attention_heads)
        self.num_key_value_heads=int(num_key_value_heads)
        self.head_dim=int(head_dim)
        self.lora_rank=int(lora_rank)
        self.hyper_lanes=int(hyper_lanes)
        self.max_position_embeddings=int(max_position_embeddings)
        self.max_loops=int(max_loops)
        self.rope_theta=float(rope_theta)
        self.rms_norm_eps=float(rms_norm_eps)
        self.ddl_beta_init=float(ddl_beta_init)
        self.ddl_k_eps=float(ddl_k_eps)
        self.ddl_v_sigmoid_scale=float(ddl_v_sigmoid_scale)
        self.refinement_cycles=int(refinement_cycles)
        self.use_cache=bool(use_cache)