from transformers import PretrainedConfig


class TransWaveConfig(PretrainedConfig):

    model_type = "transwave_meyer"


    def __init__(
        self,
        vocab_size=256,
        byte_dim=96,
        bytes_per_block=8,
        num_scales=8,
        hidden_size=768,
        max_position_embeddings=32768,
        block_context=4096,
        num_hidden_layers=12,
        num_groups=12,
        meyer_bands=12,
        meyer_controls_per_band=8,
        meyer_kernel_blocks=512,
        local_kernel_blocks=7,
        intermediate_size=2816,
        decoder_hidden_size=256,
        trained_context_length=20480,
        architecture_id="causal_transwave_blocklifting_meyer_byte_lm_g4_v2",
        **kwargs,
    ):

        # config.json may already supply this.
        # Never provide it twice.

        kwargs.setdefault(
            "tie_word_embeddings",
            True,
        )

        super().__init__(
            **kwargs
        )

        self.vocab_size = vocab_size

        self.byte_dim = byte_dim

        self.bytes_per_block = bytes_per_block

        self.num_scales = num_scales

        self.hidden_size = hidden_size

        self.max_position_embeddings = (
            max_position_embeddings
        )

        self.block_context = block_context

        self.num_hidden_layers = (
            num_hidden_layers
        )

        self.num_groups = num_groups

        self.meyer_bands = meyer_bands

        self.meyer_controls_per_band = (
            meyer_controls_per_band
        )

        self.meyer_kernel_blocks = (
            meyer_kernel_blocks
        )

        self.local_kernel_blocks = (
            local_kernel_blocks
        )

        self.intermediate_size = (
            intermediate_size
        )

        self.decoder_hidden_size = (
            decoder_hidden_size
        )

        self.trained_context_length = (
            trained_context_length
        )

        self.architecture_id = (
            architecture_id
        )