"""HuggingFace Transformers config for ExpIvme-DiffusionConversate-v1."""

from transformers import PretrainedConfig


class ExpIvmeDiffusionConfig(PretrainedConfig):
    model_type = "expivme_diffusion"

    def __init__(self, vocab_size=16001, hidden_dim=896, n_layers=12, n_heads=14,
                 context_len=1024, ffn_mult=4.0, rope_theta=10_000.0, norm_eps=1e-5,
                 tie_embeddings=True, dropout=0.0, mask_token_id=16000, **kwargs):
        self.vocab_size = vocab_size
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.context_len = context_len
        self.ffn_mult = ffn_mult
        self.rope_theta = rope_theta
        self.norm_eps = norm_eps
        self.dropout = dropout
        self.mask_token_id = mask_token_id
        self.max_position_embeddings = context_len
        self.num_hidden_layers = n_layers
        self.num_attention_heads = n_heads
        self.hidden_size = hidden_dim
        kwargs.setdefault("tie_word_embeddings", tie_embeddings)
        super().__init__(**kwargs)

    @property
    def head_dim(self):
        return self.hidden_dim // self.n_heads