from transformers import PretrainedConfig


class HaloSConfig(PretrainedConfig):
    model_type = "halo_s"

    def __init__(
        self,
        vocab_size=50257,
        hidden_size=512,
        num_layers=6,
        num_heads=8,
        num_kv_heads=2,
        num_globals=2,
        local_window=64,
        dilated_offsets=None,
        num_random=2,
        dropout=0.1,
        max_seq_len=4096,
        **kwargs
    ):
        super().__init__(**kwargs)

        if dilated_offsets is None:
            dilated_offsets = [1, 2, 4, 8, 16, 32, 64, 128]

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads

        self.num_globals = num_globals
        self.local_window = local_window
        self.dilated_offsets = dilated_offsets
        self.num_random = num_random

        self.dropout = dropout
        self.max_seq_len = max_seq_len

    @property
    def head_dim(self):
        return self.hidden_size // self.num_heads