from transformers import PretrainedConfig

class HCAConfig(PretrainedConfig):
    model_type = "hca"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size=8192,
        d_model=512,
        n_layers=16,
        n_heads=8,
        d_ff=1536,
        chunk_size=32,
        s_landmarks=4,
        seq_len=256,
        pad_token_id=0,
        bos_token_id=4,
        eos_token_id=4,
        tie_word_embeddings=False,
        **kwargs
    ):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.hidden_size = d_model
        self.n_layers = n_layers
        self.num_hidden_layers = n_layers
        self.n_heads = n_heads
        self.num_attention_heads = n_heads
        self.d_ff = d_ff
        self.intermediate_size = d_ff
        self.chunk_size = chunk_size
        self.s_landmarks = s_landmarks
        self.seq_len = seq_len
        self.max_position_embeddings = seq_len
        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,
            **kwargs
        )