"""gLM2 model configuration."""

from typing import Optional

from transformers import PretrainedConfig
from transformers.utils import logging

logger = logging.get_logger(__name__)


class gLM2Config(PretrainedConfig):
    """Configuration class for gLM2.

    Mirrors the hyperparameters in the official `tattabio/gLM2_*` checkpoints
    plus the standard HuggingFace `PretrainedConfig` knobs.
    """

    model_type = "gLM2"

    def __init__(
        self,
        dim: int = 640,
        depth: int = 30,
        heads: int = 10,
        vocab_size: int = 37,
        swiglu_multiple_of: int = 256,
        ffn_dim_multiplier: Optional[float] = None,
        norm_eps: float = 1e-5,
        pad_token_id: int = 1,
        cls_token_id: int = 0,
        eos_token_id: int = 2,
        mask_token_id: int = 35,
        max_position_embeddings: int = 4096,
        initializer_range: float = 0.02,
        tie_word_embeddings: bool = False,
        **kwargs,
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )
        self.dim = dim
        self.depth = depth
        self.heads = heads
        self.vocab_size = vocab_size
        self.swiglu_multiple_of = swiglu_multiple_of
        self.ffn_dim_multiplier = ffn_dim_multiplier
        self.norm_eps = norm_eps
        self.cls_token_id = cls_token_id
        self.eos_token_id = eos_token_id
        self.mask_token_id = mask_token_id
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range

        self.auto_map = {
            "AutoConfig": "configuration_glm2.gLM2Config",
            "AutoModel": "modeling_glm2.gLM2Model",
            "AutoModelForMaskedLM": "modeling_glm2.gLM2ForMaskedLM",
        }