from dataclasses import dataclass, field
from enum import IntEnum

# Common keys

layer_keys_llama_norms = [["input_layernorm"],
                          ["post_attention_layernorm"]]
layer_keys_cohere_norms = [["input_layernorm"]]
layer_keys_gpt2_norms = [["ln_1"],
                         ["ln_2"]]
layer_keys_yi_norms = [["ln1", "input_layernorm"],
                       ["ln2", "post_attention_layernorm"]]
layer_keys_gemma2_norms = [["input_layernorm"],
                           ["post_attention_layernorm"],
                           ["pre_feedforward_layernorm"],
                           ["post_feedforward_layernorm"]]
layer_keys_internlm2_norms = [["attention_norm"],
                              ["ffn_norm"]]
layer_keys_glm4_norms = [["input_layernorm"],
                          ["post_self_attn_layernorm"],
                          ["post_attention_layernorm"],
                          ["post_mlp_layernorm"]]
layer_keys_llama_attn = [["self_attn.q_proj"],
                         ["self_attn.k_proj"],
                         ["self_attn.v_proj"],
                         ["self_attn.o_proj"]]
layer_keys_gpt2_attn = [["self_attn.c_attn", "self_attn.q_proj"],
                        ["self_attn.c_attn", "self_attn.k_proj"],
                        ["self_attn.c_attn", "self_attn.v_proj"],
                        ["self_attn.o_proj"]]
layer_keys_internlm2_attn = [["self_attn.wqkv", "self_attn.q_proj"],
                             ["self_attn.wqkv", "self_attn.k_proj"],
                             ["self_attn.wqkv", "self_attn.v_proj"],
                             ["self_attn.o_proj"]]
layer_keys_dbrx_attn = [["self_attn.Wqkv", "self_attn.q_proj"],
                        ["self_attn.Wqkv", "self_attn.k_proj"],
                        ["self_attn.Wqkv", "self_attn.v_proj"],
                        ["self_attn.o_proj"]]
layer_keys_phi3_attn = [["self_attn.qkv_proj", "self_attn.q_proj"],
                        ["self_attn.qkv_proj", "self_attn.k_proj"],
                        ["self_attn.qkv_proj", "self_attn.v_proj"],
                        ["self_attn.o_proj"]]
layer_keys_llama_mlp = [["mlp.down_proj"],
                        ["mlp.gate_proj"],
                        ["mlp.up_proj"]]
layer_keys_internlm2_mlp = [["feed_forward.w1"],
                            ["feed_forward.w2"],
                            ["feed_forward.w3"]]
layer_keys_phi3_mlp = [["mlp.down_proj"],
                       ["mlp.gate_up_proj", "mlp.gate_proj"],
                       ["mlp.gate_up_proj", "mlp.up_proj"]]
layer_keys_mixtral_mlp = [["block_sparse_moe.experts.*.w1"],
                          ["block_sparse_moe.experts.*.w2"],
                          ["block_sparse_moe.experts.*.w3"],
                          ["block_sparse_moe.gate"]]
layer_keys_qwen3moe_mlp = [["mlp.experts.*.gate_proj"],
                           ["mlp.experts.*.up_proj"],
                           ["mlp.experts.*.down_proj"],
                           ["mlp.gate"]]
layer_keys_dbrx_mlp = [["block_sparse_moe.experts.*.v1", "block_sparse_moe.experts.v1"],
                       ["block_sparse_moe.experts.*.w1", "block_sparse_moe.experts.w1"],
                       ["block_sparse_moe.experts.*.w2", "block_sparse_moe.experts.w2"],
                       ["block_sparse_moe.gate"]]
layer_keys_llama_mlp_swiglu = [["mlp.swiglu.w12"],
                               ["mlp.swiglu.w3"]]
layer_keys_starcoder2_mlp = [["mlp.c_fc"],
                             ["mlp.c_proj"]]
layer_keys_gpt2_mlp = [["mlp.c_fc"],
                       ["mlp.c_proj"]]
expect_keys_llama = [["lm_head"],
                     ["model.norm"],
                     ["model.embed_tokens"]]
expect_keys_gemma = [["model.norm"],
                     ["model.embed_tokens"]]
expect_keys_starcoder2 = [["model.norm"],
                          ["model.embed_tokens"]]
expect_keys_gpt2 = [["model.embed_tokens"]]

dbrx_keymap = [("transformer.", "model."),
               (".blocks.", ".layers."),
               (".ffn.experts.mlp.", ".block_sparse_moe.experts."),
               (".ffn.router.layer.", ".block_sparse_moe.gate."),
               (".norm_attn_norm.norm_1.", ".input_layernorm."),
               (".norm_attn_norm.norm_2.", ".post_attention_layernorm."),
               (".norm_attn_norm.attn.", ".self_attn."),
               (".out_proj.", ".o_proj."),
               (".norm_f.", ".norm."),
               (".wte.", ".embed_tokens.")]
bigcode_keymap = [("transformer.ln_f", "model.norm"),
                  ("transformer.", "model."),
                  (".attn.c_proj.", ".self_attn.o_proj."),
                  (".attn.", ".self_attn."),
                  (".h.", ".layers."),
                  (".wte.", ".embed_tokens.")]
gpt2_keymap = [("$ln_f.", "model.norm."),
               (".attn.c_proj.", ".self_attn.o_proj."),
               (".attn.", ".self_attn."),
               ("$h.", "model.layers."),
               ("$wte.", "model.embed_tokens."),
               ("$wpe.", "model.wpe.")]
internlm2_keymap = [("$output.", "lm_head."),
                    ("$model.tok_embeddings.", "model.embed_tokens."),
                    (".attention.", ".self_attn."),
                    (".wo.", ".o_proj.")]
google_keymap = [("mm_input_projection_weight", "mm_input_projection.weight")]

no_default = object()

class RopeStyle(IntEnum):
    NONE = 0
    GPTJ = 1
    NEOX = 2

class ExLlamaV2ArchParams:

    def __init__(self, arch_string: str, read_config: dict):
        """
        Get architecture definition from model config. If the architecture isn't recognized, defaults to Llama
        architecture.

        :param arch_string:
            Architecture string from config.json

        :param read_config:
            config.json as Python dict
        """

        self.arch_string = arch_string
        arch_recognized = False

        self.keymap = None
        self.compile_fix_keymap = None

        @dataclass
        class Params:
            keys: dict = field(default_factory = lambda: {
                "norm_eps": "rms_norm_eps",
                "norm_1": ".input_layernorm",
                "norm_1_post": None,
                "fused_qkv": None,
                "mlp_gate": ".mlp.gate_proj",
                "mlp_up": ".mlp.up_proj",
                "mlp_down": ".mlp.down_proj",
                "lm_head": "lm_head",
                "norm_2": ".post_attention_layernorm",
                "norm_2_post": None,
                "fused_mlp_12": None,
                "fused_mlp_3": None,
                "learned_pos_emb": None,
                "attn_q": ".self_attn.q_proj",
                "attn_k": ".self_attn.k_proj",
                "attn_v": ".self_attn.v_proj",
                "attn_o": ".self_attn.o_proj",
                "layers": "layers",
                "patch_conv": "patch_conv",
            })

            # Compute logit scale from `dim_model_base` key in config.json (MiniCPM quirk)
            logit_scale_basedim = False

            # Clamp hidden states to FP16 range
            clamp_hidden_states = False

            # Upcast hidden state to FP32 before adding to residual stream
            residual_stream_fp32 = False

            # Normalize embeddings (Gemma quirk)
            normalize_embeddings = False

            # Constant bias for layernorm (Gemma quirk)
            norm_constant_bias = 0

            # Alternate packing scheme for fused QKV tensor (InternLM2 quirk)
            fused_qkv_altpack = False

            # SWA required by architecture
            swa = False
            alternating_swa = False
            sliding_rope_theta = None
            sliding_rope_scale = None
            pos_id_index = 0

            # Model only works with eager attention
            eager_attn_only = False

            # Expect bias for linear layers
            attention_bias_qkv = False
            attention_bias_o = False
            mlp_bias = False

            # Default multiplier for MLP inner dim (GPT2 quirk)
            default_inner_dim_mult = None

            # Use gated MLP
            mlp_gate = True

            # Use block-sparse MLP
            is_moe = False

            # Use parallel decoder blocks (Cohere quirk)
            parallel_decoder_blocks = False

            # Use MQA, effectively num_key_value_heads = 1 (GPTBigCode quirk)
            mqa = False

            # Model is incoherent without BOS at the start of the context
            requires_bos = False

            # Scale attn weights (GPT2 quirk, not important for inference)
            scale_attn_weights = False

            # Model implementation works in tensor-parallel mode
            supports_tp = False

            # Activation function
            mlp_act_func = "silu"

            # Layer norm type
            norm = "rmsnorm"
            headnorm = "layernorm"

            # RoPE style
            rope_style = RopeStyle.NEOX

            # Expected keys
            expect_keys: list[str] = field(default_factory = lambda: [])
            layer_keys: list[str] = field(default_factory = lambda: [])

            # Defaults because Gemma3
            default_vocab_size = no_default
            default_rms_norm_eps = no_default
            default_head_dim = no_default
            default_num_attention_heads = no_default
            default_num_key_value_heads = no_default
            default_use_qk_norm = False
            default_sliding_window_pattern = 1
            default_rope_theta = 10000

            # Vision stuff
            patch_conv_bias: bool = False
            is_vision: bool = False
            vision_input_norm: bool = True
            vision_conv3d: bool = False
            mrope: bool = False
            rope_freq_half: bool = False
            learned_emb: bool = False
            output_norm: bool = False
            mlp_merger: bool = False
            mlp_patch_merger: bool = False

        # Component models
        self.lm_prefix = ""
        self.vt_prefix = ""
        self.mmp_prefix = ""
        self.lm = Params()
        self.mmp = Params()
        self.vt = Params()

        self.mmp.keys.update({
            "norm_1": None,
            "norm_1_post": None,
            "norm_2": None,
            "norm_2_post": None,
            "fused_mlp_12": None,
            "fused_mlp_3": None,
        })
        self.mmp.rope_style = RopeStyle.NONE

        self.vt.is_vision = True

        # Tensors are transposed in original model weights
        self.orig_weights_transposed = False

        # Add noise rows to calibration while quantizing
        self.standard_calib_noise = None

        # Mistral

        if arch_string == "MistralForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.supports_tp = True

        # Mixtral

        if arch_string == "MixtralForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_mixtral_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.keys.update({
                "mlp_gate": ".block_sparse_moe.experts.*.w1",
                "mlp_up": ".block_sparse_moe.experts.*.w3",
                "mlp_down": ".block_sparse_moe.experts.*.w2",
                "mlp_expert_gate": ".block_sparse_moe.gate"
            })
            self.lm.is_moe = True

        # Pixtral

        if (
            arch_string == "LlavaForConditionalGeneration" and
            "vision_config" in read_config and
            read_config["vision_config"].get("model_type") == "pixtral"
        ):
            arch_recognized = True
            self.lm_prefix = "language_model."
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama

            self.vt_prefix = "vision_tower."
            self.vt.keys.update({
                "attn_q": ".attention.q_proj",
                "attn_k": ".attention.k_proj",
                "attn_v": ".attention.v_proj",
                "attn_o": ".attention.o_proj",
                "mlp_gate": ".feed_forward.gate_proj",
                "mlp_up": ".feed_forward.up_proj",
                "mlp_down": ".feed_forward.down_proj",
                "norm_1": ".attention_norm",
                "norm_2": ".ffn_norm",
                "layers": "transformer.layers",
                "ln_pre": "ln_pre",
            })
            self.vt.mlp_merger = True

            self.mmp_prefix = "multi_modal_projector."
            self.mmp.keys.update({
                "mlp_gate": None,
                "mlp_up": "linear_1",
                "mlp_down": "linear_2",
            })
            self.mmp.mlp_gate = False
            self.mmp.mlp_act_func = "gelu"
            self.mmp.mlp_bias = bool(read_config.get("multimodal_projector_bias", True))

        # Mistral 3 multimodal

        if (
            arch_string == "Mistral3ForConditionalGeneration" and
            "vision_config" in read_config and
            read_config["vision_config"].get("model_type") == "pixtral"
        ):
            arch_recognized = True
            self.lm_prefix = "language_model."
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama

            self.vt_prefix = "vision_tower."
            self.vt.keys.update({
                "attn_q": ".attention.q_proj",
                "attn_k": ".attention.k_proj",
                "attn_v": ".attention.v_proj",
                "attn_o": ".attention.o_proj",
                "mlp_gate": ".feed_forward.gate_proj",
                "mlp_up": ".feed_forward.up_proj",
                "mlp_down": ".feed_forward.down_proj",
                "norm_1": ".attention_norm",
                "norm_2": ".ffn_norm",
                "layers": "transformer.layers",
                "ln_pre": "ln_pre",
            })
            self.vt.mlp_merger = True
            self.vt.mlp_patch_merger = True

            self.mmp_prefix = "multi_modal_projector."
            self.mmp.keys.update({
                "norm_2": "norm",
                "mlp_gate": None,
                "mlp_up": "linear_1",
                "mlp_down": "linear_2",
                "patch_merger": "patch_merger.merging_layer",
            })
            self.mmp.mlp_patch_merger = True
            self.mmp.mlp_gate = False
            self.mmp.mlp_act_func = "gelu"
            self.mmp.mlp_bias = bool(read_config.get("multimodal_projector_bias", True))

        # Yi

        if arch_string == "YiForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_yi_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.keys.update({
                "norm_1": ".ln1",
                "norm_2": ".ln2",
            })

        # Orion

        if arch_string == "OrionForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.norm = "layernorm"

        # Qwen2 (1.5, 2, 2.5)

        if arch_string == "Qwen2ForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.attention_bias_qkv = True
            self.lm.supports_tp = True

        # Qwen3

        if arch_string == "Qwen3ForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.supports_tp = True
            self.lm.default_use_qk_norm = True

        # Qwen3MoE

        if arch_string == "Qwen3MoeForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_qwen3moe_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.supports_tp = True
            self.lm.default_use_qk_norm = True
            self.lm.keys.update({
                "mlp_gate": ".mlp.experts.*.gate_proj",
                "mlp_up": ".mlp.experts.*.up_proj",
                "mlp_down": ".mlp.experts.*.down_proj",
                "mlp_expert_gate": ".mlp.gate"
            })
            self.lm.is_moe = True

        # Qwen2-VL (2, 2.5)

        if arch_string in ["Qwen2VLForConditionalGeneration", "Qwen2_5_VLForConditionalGeneration"]:
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.attention_bias_qkv = True
            self.lm.mrope = True
            self.lm.rope_freq_half = True

            self.vt_prefix = "visual."
            if arch_string == "Qwen2VLForConditionalGeneration":
                read_config["vision_config"].update({"model_type": "qwen2"})
                self.vt.keys.update({
                    "fused_qkv": ".attn.qkv",
                    "attn_o": ".attn.proj",
                    "mlp_gate": None,
                    "mlp_up": ".mlp.fc1",
                    "mlp_down": ".mlp.fc2",
                    "norm_1": ".norm1",
                    "norm_2": ".norm2",
                    "layers": "blocks",
                    "patch_conv": "patch_embed.proj",
                })
                self.vt.mlp_gate = False
                self.vt.mlp_act_func = "quickgelu"
                self.vt.norm = "layernorm"
            elif arch_string == "Qwen2_5_VLForConditionalGeneration":
                read_config["vision_config"].update({"model_type": "qwen2.5"})
                self.vt.keys.update({
                    "fused_qkv": ".attn.qkv",
                    "attn_o": ".attn.proj",
                    "mlp_gate": ".mlp.gate_proj",
                    "mlp_up": ".mlp.up_proj",
                    "mlp_down": ".mlp.down_proj",
                    "norm_1": ".norm1",
                    "norm_2": ".norm2",
                    "layers": "blocks",
                    "patch_conv": "patch_embed.proj",
                })
                self.vt.mlp_gate = True
                self.vt.mlp_act_func = "silu"
                self.vt.norm = "rmsnorm"
            self.vt.mlp_bias = True
            self.vt.attention_bias_qkv = True
            self.vt.attention_bias_o = True
            self.vt.vision_input_norm = False
            self.vt.vision_conv3d = True
            self.vt.mlp_merger = True

            self.mmp_prefix = "visual.merger."
            self.mmp.keys.update({
                "mlp_gate": None,
                "mlp_up": "mlp.0",
                "mlp_down": "mlp.2",
                "norm_2": "ln_q",
            })
            self.mmp.mlp_gate = False
            self.mmp.mlp_act_func = "gelu"
            self.mmp.mlp_bias = True
            self.mmp.norm = "layernorm"

            self.standard_calib_noise = (5, 30)

        # Gemma

        if arch_string == "GemmaForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_gemma
            self.lm.keys.update({
                "lm_head": "model.embed_tokens",
            })
            self.lm.mlp_act_func = "gelu"
            self.lm.normalize_embeddings = True
            self.lm.norm_constant_bias = 1
            self.lm.requires_bos = True

        # Gemma2

        if arch_string == "Gemma2ForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_gemma2_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_gemma
            self.lm.keys.update({
                "lm_head": "model.embed_tokens",
                "norm_1": ".input_layernorm",
                "norm_1_post": ".post_attention_layernorm",
                "norm_2": ".pre_feedforward_layernorm",
                "norm_2_post": ".post_feedforward_layernorm",
            })
            self.lm.mlp_act_func = "gelu"
            self.lm.normalize_embeddings = True
            self.lm.norm_constant_bias = 1
            self.lm.requires_bos = True
            self.lm.alternating_swa = True
            self.lm.residual_stream_fp32 = True

        # Gemma3

        if arch_string == "Gemma3ForConditionalGeneration":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_gemma2_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_gemma
            self.lm.keys.update({
                "lm_head": "model.embed_tokens",
                "norm_1": ".input_layernorm",
                "norm_1_post": ".post_attention_layernorm",
                "norm_2": ".pre_feedforward_layernorm",
                "norm_2_post": ".post_feedforward_layernorm",
            })
            self.lm_prefix = "language_model."
            self.lm.mlp_act_func = "gelu"
            self.lm.normalize_embeddings = True
            self.lm.norm_constant_bias = 1
            self.lm.requires_bos = True
            self.lm.alternating_swa = True
            self.lm.residual_stream_fp32 = True
            self.lm.sliding_rope_theta = 10000
            self.lm.sliding_rope_scale = 1
            self.lm.default_vocab_size = 262208
            self.lm.default_rms_norm_eps = 1e-06
            self.lm.default_head_dim = 256
            self.lm.default_num_attention_heads = 8
            self.lm.default_num_key_value_heads = 4
            self.lm.default_use_qk_norm = True
            self.lm.default_sliding_window_pattern = 6
            self.lm.default_rope_theta = 1e6
            self.lm.pos_id_index = 1
            self.lm.headnorm = "rmsnorm"

            self.vt_prefix = "vision_tower.vision_model."
            self.vt.keys.update({
                "attn_q": ".self_attn.q_proj",
                "attn_k": ".self_attn.k_proj",
                "attn_v": ".self_attn.v_proj",
                "attn_o": ".self_attn.out_proj",
                "norm_1": ".layer_norm1",
                "norm_2": ".layer_norm2",
                "mlp_gate": None,
                "mlp_up": ".mlp.fc1",
                "mlp_down": ".mlp.fc2",
                "layers": "encoder.layers",
                "patch_conv": "embeddings.patch_embedding",
                "position_embedding": "embeddings.position_embedding",
                "output_norm": "post_layernorm",
            })
            self.vt.norm = "rmsnorm"
            self.vt.patch_conv_bias = True
            self.vt.mlp_gate = False
            self.vt.mlp_bias = True
            self.vt.attention_bias_qkv = True
            self.vt.attention_bias_o = True
            self.vt.vision_input_norm = False
            self.vt.mlp_merger = False
            self.vt.norm = "layernorm"
            self.vt.learned_emb = True
            self.vt.rope_style = RopeStyle.NONE
            self.vt.mlp_act_func = "gelu"
            self.vt.output_norm = True

            self.keymap = google_keymap
            self.compile_fix_keymap = google_keymap
            self.mmp_prefix = "multi_modal_projector."
            self.mmp.keys.update({
                "input_projection": "mm_input_projection",
                "input_projection_norm": "mm_soft_emb_norm",
            })
            self.mmp.norm_constant_bias = 1

        # StarCoder2

        if arch_string == "Starcoder2ForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_starcoder2_mlp
            self.lm.expect_keys += \
                expect_keys_starcoder2
            self.lm.keys.update({
                "mlp_gate": None,
                "mlp_up": ".mlp.c_fc",
                "mlp_down": ".mlp.c_proj",
                "lm_head": "model.embed_tokens",
                "norm_eps": "layer_norm_epsilon",
            })
            self.lm.mlp_act_func = "gelu"
            self.lm.norm = "layernorm"
            self.lm.attention_bias_qkv = True
            self.lm.attention_bias_o = True
            self.lm.mlp_bias = True
            self.lm.mlp_gate = False

        # GemMoE

        if arch_string == "GemmoeForCausalLM":
            arch_recognized = True
            print(f" !! Warning, Gemmoe support is experimental and has not been fully tested")
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_mixtral_mlp
            self.lm.expect_keys += \
                expect_keys_gemma
            self.lm.keys.update({
                "mlp_gate": ".block_sparse_moe.experts.*.w1",
                "mlp_up": ".block_sparse_moe.experts.*.w3",
                "mlp_down": ".block_sparse_moe.experts.*.w2",
                "mlp_expert_gate": ".block_sparse_moe.gate",
                "lm_head": "model.embed_tokens",
            })
            self.lm.mlp_act_func = "gelu"
            self.lm.normalize_embeddings = True
            self.lm.norm_constant_bias = 1
            self.lm.is_moe = True
            self.lm.requires_bos = True

        # Cohere

        if arch_string == "CohereForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_cohere_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_gemma
            self.lm.keys.update({
                "norm_eps": "layer_norm_eps",
                "lm_head": "model.embed_tokens",
                "norm_1": ".input_layernorm",
                "norm_2": None,
            })
            self.lm.norm = "layernorm"
            self.lm.rope_style = RopeStyle.GPTJ
            self.lm.parallel_decoder_blocks = True
            self.lm.requires_bos = True

        # Cohere 2

        if arch_string == "Cohere2ForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_cohere_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_gemma
            self.lm.keys.update({
                "norm_eps": "layer_norm_eps",
                "lm_head": "model.embed_tokens",
                "norm_1": ".input_layernorm",
                "norm_2": None,
            })
            self.lm.norm = "layernorm"
            self.lm.rope_style = RopeStyle.GPTJ
            self.lm.parallel_decoder_blocks = True
            self.lm.requires_bos = True
            self.lm.alternating_swa = True

        # DBRX

        if arch_string == "DbrxForCausalLM":
            arch_recognized = True
            self.keymap = dbrx_keymap
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_dbrx_attn + \
                layer_keys_dbrx_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.keys.update({
                "norm_eps": None,
                "mlp_gate": ".block_sparse_moe.experts.*.w1",
                "mlp_up": ".block_sparse_moe.experts.*.v1",
                "mlp_down": ".block_sparse_moe.experts.*.w2",
                "mlp_expert_gate": ".block_sparse_moe.gate",
                "fused_qkv": ".self_attn.Wqkv",
            })
            self.lm.norm = "layernorm"
            self.lm.is_moe = True

        # Phi3

        if arch_string == "Phi3ForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_phi3_attn + \
                layer_keys_phi3_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.keys.update({
                "fused_qkv": ".self_attn.qkv_proj",
                "fused_mlp_12": "gate_up_proj",
            })

        # GPTBigCode

        if arch_string == "GPTBigCodeForCausalLM":
            arch_recognized = True
            self.keymap = bigcode_keymap
            self.lm.layer_keys += \
                layer_keys_gpt2_norms + \
                layer_keys_gpt2_attn + \
                layer_keys_gpt2_mlp
            self.lm.expect_keys += \
                expect_keys_gpt2
            self.lm.keys.update({
                "norm_eps": "layer_norm_epsilon",
                "mlp_gate": None,
                "mlp_up": ".mlp.c_fc",
                "mlp_down": ".mlp.c_proj",
                "lm_head": "model.embed_tokens",
                "norm_1": ".ln_1",
                "norm_2": ".ln_2",
                "fused_qkv": ".self_attn.c_attn",
                "learned_pos_emb": "model.wpe",
            })
            self.lm.mlp_act_func = "gelu"
            self.lm.norm = "layernorm"
            self.lm.rope_style = RopeStyle.NONE
            self.lm.mqa = True
            self.lm.attention_bias_qkv = True
            self.lm.attention_bias_o = True
            self.lm.mlp_bias = True
            self.lm.mlp_gate = False

        # GPT2

        if arch_string == "GPT2LMHeadModel":
            arch_recognized = True
            self.keymap = gpt2_keymap
            self.lm.layer_keys += \
                layer_keys_gpt2_norms + \
                layer_keys_gpt2_attn + \
                layer_keys_gpt2_mlp
            self.lm.expect_keys += \
                expect_keys_gpt2
            self.lm.keys.update({
                "norm_eps": "layer_norm_epsilon",
                "mlp_gate": None,
                "mlp_up": ".mlp.c_fc",
                "mlp_down": ".mlp.c_proj",
                "lm_head": "model.embed_tokens",
                "norm_1": ".ln_1",
                "norm_2": ".ln_2",
                "fused_qkv": ".self_attn.c_attn",
                "learned_pos_emb": "model.wpe",
            })
            self.lm.mlp_act_func = "gelu"
            self.lm.norm = "layernorm"
            self.lm.rope_style = RopeStyle.NONE
            self.lm.default_inner_dim_mult = 4
            self.lm.attention_bias_qkv = True
            self.lm.attention_bias_o = True
            self.lm.mlp_bias = True
            self.lm.mlp_gate = False
            self.orig_weights_transposed = True

        # MiniCPM

        if arch_string == "MiniCPMForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.logit_scale_basedim = True

        # InternLM2

        if arch_string == "InternLM2ForCausalLM":
            arch_recognized = True
            self.keymap = internlm2_keymap
            self.lm.layer_keys += \
                layer_keys_internlm2_norms + \
                layer_keys_internlm2_attn + \
                layer_keys_internlm2_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.keys.update({
                "mlp_gate": ".feed_forward.w1",
                "mlp_up": ".feed_forward.w3",
                "mlp_down": ".feed_forward.w2",
                "norm_1": ".attention_norm",
                "norm_2": ".ffn_norm",
                "fused_qkv": ".self_attn.wqkv",
            })
            self.lm.fused_qkv_altpack = True

        # Index

        if arch_string == "IndexForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama

        # Granite (v3)

        if arch_string == "GraniteForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama

        # GLM4

        if arch_string == "Glm4ForCausalLM":
            arch_recognized = True
            self.lm.layer_keys += \
                layer_keys_glm4_norms + \
                layer_keys_llama_attn + \
                layer_keys_phi3_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.supports_tp = True
            self.lm.rope_style = RopeStyle.GPTJ
            self.lm.keys.update({
                "fused_mlp_12": "gate_up_proj",
                "lm_head": "model.embed_tokens",
                "norm_1": ".input_layernorm",
                "norm_1_post": ".post_self_attn_layernorm",
                "norm_2": ".post_attention_layernorm",
                "norm_2_post": ".post_mlp_layernorm",
            })
            self.lm.attention_bias_qkv = read_config.get("attention_bias", False)

        # Llama (default + fallback)

        if arch_string != "LlamaForCausalLM" and not arch_recognized:
            print(f" !! Warning, unknown architecture: {arch_string}")
            print(f" !! Loading as LlamaForCausalLM")
            self.arch_string = "LlamaForCausalLM"
        if not arch_recognized:
            self.lm.layer_keys += \
                layer_keys_llama_norms + \
                layer_keys_llama_attn + \
                layer_keys_llama_mlp
            self.lm.expect_keys += \
                expect_keys_llama
            self.lm.supports_tp = True

        # Arch overrides

        if read_config.get("attention_bias", False) and not (self.lm.attention_bias_qkv or self.lm.attention_bias_o):
            self.lm.attention_bias_qkv = True
            self.lm.attention_bias_o = True

        if read_config.get("mlp_bias", False):
            self.lm.mlp_bias = True

        if read_config.get("tie_word_embeddings", False):
            if ["lm_head"] in self.lm.expect_keys:
                self.lm.expect_keys.remove(["lm_head"])
                self.lm.keys.update({
                    "lm_head": "model.embed_tokens",
                })

        # Sanity checks

        if self.lm.residual_stream_fp32:
            assert self.lm.keys["norm_1_post"] and self.lm.keys["norm_2_post"], \
                "FP32 residual stream only implemented for arch with post layernorms"

    def make_fused_mlp(self):

        for x in layer_keys_llama_mlp: self.lm.layer_keys.remove(x)
        self.lm.layer_keys += layer_keys_llama_mlp_swiglu
        self.lm.keys.update({
            "fused_mlp_12": layer_keys_llama_mlp_swiglu[0][0],
            "fused_mlp_3": layer_keys_llama_mlp_swiglu[1][0],
        })
