# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
# SPDX-License-Identifier: Apache-2.0
# Contact: qubitium@modelcloud.ai, x.com/qubitium
from transformers import AutoModelForImageTextToText

from . import LlamaQModel


class Qwen3_5QModel(LlamaQModel):
    """
    Qwen3_5 inherits the Llama-style layout but inserts Q/K RMS norm layers
    ahead of the attention projections. We mark those helper modules as
    non-quantized so the layer walker captures the complete structure.
    """

    loader = AutoModelForImageTextToText

    require_load_processor = True

    layer_modules_strict = False

    pre_lm_head_norm_module = "model.language_model.norm"

    rotary_embedding = "model.language_model.rotary_emb"

    # Preserve auxiliary MTP/draft-head tensors when present.
    # Qwen3_5_MoeQModel already does this; dense Qwen3.5/Qwen3.6 models
    # can also ship mtp.* tensors in auxiliary safetensors files.
    out_of_model_tensors = {"prefixes": ["mtp"]}

    module_tree = [
        "model",
        "language_model",
        "layers",
        "#",
        {
            "input_layernorm": ("input_layernorm:!",),
            "self_attn": ("q_norm:!", "q_proj:0:in=x", "k_norm:!", "k_proj:0:in=x", "v_proj:0:in=x", "o_proj:1"),
            "linear_attn": (
                "norm:!",
                "conv1d:!",
                "in_proj_qkv:0:in=x",
                "in_proj_z:1:in=x",
                "in_proj_b:!:1",
                "in_proj_a:!:1",
                "out_proj:2",
            ),
            "post_attention_layernorm": ("post_attention_layernorm:!",),
            "mlp": ("gate_proj:0:in=x", "up_proj:0:in=x", "down_proj:1"),
        },
    ]