"""Configuration of the Sansar Sanskrit language models (Muse Mesh).

Code licence: Apache-2.0 (LICENSE-CODE). The model weights are licensed separately (LICENSE).
"""
from transformers import PretrainedConfig


class SansarConfig(PretrainedConfig):
    """A compact pre-LayerNorm GPT. The flags select the two architectures the Sansar runs used:

    - ``rope=False, qk_norm=False, untied_head=False, relu_sq=False``: learned position table, GELU (tanh) MLP, output
      head tied to the token embedding (the E8 to F3 models).
    - ``rope=True, qk_norm=True, untied_head=True, relu_sq=True``: rotary positions, RMS-normalised queries and keys,
      a separate output head, squared-ReLU MLP (the F4 recipe and every model after it).

    ``value_embed`` and ``logit_cap`` are supported for completeness; no released model uses them.
    """

    model_type = "sansar"
    attribute_map = {
        "hidden_size": "n_embd",
        "num_hidden_layers": "n_layer",
        "num_attention_heads": "n_head",
        "max_position_embeddings": "block_size",
    }

    def __init__(
        self,
        vocab_size: int = 8000,
        block_size: int = 512,
        n_layer: int = 6,
        n_head: int = 8,
        n_embd: int = 512,
        mlp_ratio: int = 4,
        bias: bool = False,
        rope: bool = False,
        rope_theta: float = 10000.0,
        qk_norm: bool = False,
        untied_head: bool = False,
        relu_sq: bool = False,
        value_embed: bool = False,
        value_embed_layers: str = "",
        logit_cap: float = 0.0,
        layer_norm_eps: float = 1e-5,
        initializer_range: float = 0.02,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        pad_token_id: int = 2,
        tie_word_embeddings=None,
        use_cache: bool = False,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.block_size = block_size
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_embd = n_embd
        self.mlp_ratio = mlp_ratio
        self.bias = bias
        self.rope = rope
        self.rope_theta = rope_theta
        self.qk_norm = qk_norm
        self.untied_head = untied_head
        self.relu_sq = relu_sq
        self.value_embed = value_embed
        self.value_embed_layers = value_embed_layers
        self.logit_cap = logit_cap
        self.layer_norm_eps = layer_norm_eps
        self.initializer_range = initializer_range
        self.use_cache = use_cache
        if tie_word_embeddings is None:
            tie_word_embeddings = not untied_head
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )