import os
import shutil
import torch
from datasets import load_dataset
from transformers import AutoTokenizer

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

if hasattr(torch, "accelerator") and not hasattr(torch.accelerator, "get_memory_info"):
    torch.accelerator.get_memory_info = lambda idx=0: torch.cuda.mem_get_info(idx)

try:
    from transformers import AutoModelForImageTextToText
    AutoModelClass = AutoModelForImageTextToText
except ImportError:
    from transformers import AutoModelForCausalLM
    AutoModelClass = AutoModelForCausalLM

from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier


MODEL_ID = "/root/autodl-tmp/Huihui-Qwen3.5-9B-abliterated"
SAVE_DIR = "/root/autodl-tmp/Huihui-Qwen3.5-9B-abliterated-AWQ-W4A16-vLLM"

DATASET_DIR = "/root/autodl-tmp/pile-val-backup"

MAX_SEQ_LENGTH = 2048
NUM_CALIBRATION_SAMPLES = 128


recipe = [
    AWQModifier(
        config_groups={
            "group_1": {
                "targets": ["Linear"],
                "weights": {
                    "num_bits": 4,
                    "type": "int",
                    "symmetric": True,
                    "group_size": 32,
                    "strategy": "group",
                    "observer": "mse",
                    "dynamic": False,
                },
                "input_activations": None,
                "output_activations": None,
            }
        },
        targets=["Linear"],
        ignore=[
            "re:.*embed_tokens",
            "re:.*linear_attn[.]in_proj_a",
            "re:.*linear_attn[.]in_proj_b",
            "re:model[.]visual.*",
            "re:mtp.*",
            "lm_head",
        ],
        mappings=[
            {
                "smooth_layer": r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]input_layernorm",
                "balance_layers": [
                    r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]q_proj",
                    r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]k_proj",
                    r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]v_proj",
                ],
            },
            {
                "smooth_layer": r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]v_proj",
                "balance_layers": [
                    r"re:model.*layers[.](3|7|11|15|19|23|27|31)[.]self_attn[.]o_proj",
                ],
            },
            {
                "smooth_layer": r"re:model.*post_attention_layernorm",
                "balance_layers": [
                    r"re:model.*mlp[.]gate_proj",
                    r"re:model.*mlp[.]up_proj",
                ],
            },
        ],
        duo_scaling=True,
        n_grid=20,
    )
]


print("Loading model...")
model = AutoModelClass.from_pretrained(
    MODEL_ID,
    dtype="auto",
    device_map="auto",
    trust_remote_code=True,
)

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

print("Loading calibration dataset...")
dataset = load_dataset(
    DATASET_DIR,
    split=f"validation[:{NUM_CALIBRATION_SAMPLES}]",
)

print("Starting AWQ quantization...")
oneshot(
    model=model,
    tokenizer=tokenizer,
    dataset=dataset,
    recipe=recipe,
    output_dir=SAVE_DIR,
    max_seq_length=MAX_SEQ_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)

# 补齐多模态 processor 文件，但不要覆盖 config.json
extra_files = [
    "preprocessor_config.json",
    "processor_config.json",
    "video_preprocessor_config.json",
    "chat_template.jinja",
    "generation_config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "vocab.json",
    "merges.txt",
]

for name in extra_files:
    src = os.path.join(MODEL_ID, name)
    dst = os.path.join(SAVE_DIR, name)
    if os.path.exists(src) and not os.path.exists(dst):
        shutil.copy2(src, dst)

print(f"Saved vLLM-compatible AWQ W4A16 model to: {SAVE_DIR}")