from ._base import BaseGPTQForCausalLM


class YuanGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "YuanDecoderLayer"
    layers_block_name = "model.layers"
    outside_layer_modules = ["model.embed_tokens", "model.norm"]
    inside_layer_modules = [
        ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
        ["self_attn.o_proj"],
        ["mlp.gate.query", "mlp.gate.key", "mlp.gate.value"],
        [
            "mlp.experts.0.w1",
            "mlp.experts.1.w1",
            "mlp.experts.2.w1",
            "mlp.experts.3.w1",
            "mlp.experts.4.w1",
            "mlp.experts.5.w1",
            "mlp.experts.6.w1",
            "mlp.experts.7.w1",
            "mlp.experts.8.w1",
            "mlp.experts.9.w1",
            "mlp.experts.10.w1",
            "mlp.experts.11.w1",
            "mlp.experts.12.w1",
            "mlp.experts.13.w1",
            "mlp.experts.14.w1",
            "mlp.experts.15.w1",
            "mlp.experts.16.w1",
            "mlp.experts.17.w1",
            "mlp.experts.18.w1",
            "mlp.experts.19.w1",
            "mlp.experts.20.w1",
            "mlp.experts.21.w1",
            "mlp.experts.22.w1",
            "mlp.experts.23.w1",
            "mlp.experts.24.w1",
            "mlp.experts.25.w1",
            "mlp.experts.26.w1",
            "mlp.experts.27.w1",
            "mlp.experts.28.w1",
            "mlp.experts.29.w1",
            "mlp.experts.30.w1",
            "mlp.experts.31.w1",
        ],
        [
            "mlp.experts.0.w2",
            "mlp.experts.1.w2",
            "mlp.experts.2.w2",
            "mlp.experts.3.w2",
            "mlp.experts.4.w2",
            "mlp.experts.5.w2",
            "mlp.experts.6.w2",
            "mlp.experts.7.w2",
            "mlp.experts.8.w2",
            "mlp.experts.9.w2",
            "mlp.experts.10.w2",
            "mlp.experts.11.w2",
            "mlp.experts.12.w2",
            "mlp.experts.13.w2",
            "mlp.experts.14.w2",
            "mlp.experts.15.w2",
            "mlp.experts.16.w2",
            "mlp.experts.17.w2",
            "mlp.experts.18.w2",
            "mlp.experts.19.w2",
            "mlp.experts.20.w2",
            "mlp.experts.21.w2",
            "mlp.experts.22.w2",
            "mlp.experts.23.w2",
            "mlp.experts.24.w2",
            "mlp.experts.25.w2",
            "mlp.experts.26.w2",
            "mlp.experts.27.w2",
            "mlp.experts.28.w2",
            "mlp.experts.29.w2",
            "mlp.experts.30.w2",
            "mlp.experts.31.w2",
        ],
    ]


__all__ = ["YuanGPTQForCausalLM"]
