"""Compression recipe used to create this checkpoint.

Run quantization from the original BF16 checkpoint, not from an already
quantized checkpoint.

AWQ needs calibration data. This build used 128 samples from
`lmms-lab/flickr30k`, each rendered at 448x448 (a single InternVL patch) with
generic English/Korean scene-description prompts. Re-run with in-domain images
if your target distribution is far from everyday photography.

Targets are regex-scoped to the language decoder, so the vision tower,
multimodal projector, input embeddings, and lm_head are never matched and stay
in BF16. `ignore` is therefore empty by design.
"""

from llmcompressor.modifiers.awq import AWQModifier
from llmcompressor.modifiers.quantization import QuantizationModifier

ATTENTION = r"re:^model\.language_model\.layers\.\d+\.self_attn\.(q_proj|k_proj|v_proj|o_proj)$"
MLP = r"re:^model\.language_model\.layers\.\d+\.mlp\.(gate_proj|up_proj|down_proj)$"

RECIPE = [
    AWQModifier(
        mappings=[
            {
                "smooth_layer": r"re:^model\.language_model\.layers\.\d+\.input_layernorm$",
                "balance_layers": [
                    r"re:^model\.language_model\.layers\.\d+\.self_attn\.q_proj$",
                    r"re:^model\.language_model\.layers\.\d+\.self_attn\.k_proj$",
                    r"re:^model\.language_model\.layers\.\d+\.self_attn\.v_proj$",
                ],
            },
            {
                "smooth_layer": r"re:^model\.language_model\.layers\.\d+\.self_attn\.v_proj$",
                "balance_layers": [
                    r"re:^model\.language_model\.layers\.\d+\.self_attn\.o_proj$",
                ],
            },
            {
                "smooth_layer": r"re:^model\.language_model\.layers\.\d+\.post_attention_layernorm$",
                "balance_layers": [
                    r"re:^model\.language_model\.layers\.\d+\.mlp\.gate_proj$",
                    r"re:^model\.language_model\.layers\.\d+\.mlp\.up_proj$",
                ],
            },
            {
                "smooth_layer": r"re:^model\.language_model\.layers\.\d+\.mlp\.up_proj$",
                "balance_layers": [
                    r"re:^model\.language_model\.layers\.\d+\.mlp\.down_proj$",
                ],
            },
        ],
        duo_scaling=True,
        n_grid=20,
    ),
    QuantizationModifier(
        targets=[ATTENTION, MLP],
        ignore=[],
        scheme="W4A16_ASYM",
    ),
]