import torch
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# 1. 路径定义
BASE_MODEL_PATH = "./Qwen3-4B-Thinking-2507"
LORA_PATH = "./Qwen3-4B-Thinking-2507_FTtrained_lora"
OUTPUT_PRUNED_PATH = "./QiMing-Polaris-Qwen3-4B-HardPruned"

# 2. 定义需要剪掉的层（根据你的 Audit 数据，L-29 到 L-32 发生了剧烈塌缩）
LAYERS_TO_REMOVE = [29, 30, 31, 32]

print("🚀 阶段 1/4: 正在加载基础模型与分词器...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH, trust_remote_code=True)
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_PATH,
    dtype=torch.bfloat16,  # 避免 deprecated 警告
    device_map="cpu",  # 融合与剪枝在 CPU 内存上进行，防止显存溢出
    trust_remote_code=True,
)

print("🔗 阶段 2/4: 加载并融合 LoRA 适配器...")
lora_model = PeftModel.from_pretrained(base_model, LORA_PATH)
merged_model = lora_model.merge_and_unload()
print("✅ LoRA 权重融合成功！")

print("✂️ 阶段 3/4: 正在执行硬剪枝 (Hard Pruning)...")
original_layers_count = len(merged_model.model.layers)
print(f"原模型总层数: {original_layers_count} 层")

# 计算需要保留的层索引
keep_indices = [i for i in range(original_layers_count) if i not in LAYERS_TO_REMOVE]
print(f"准备移除的层: {LAYERS_TO_REMOVE}")
print(f"保留的层索引数量: {len(keep_indices)} 层")

# 1. 物理切割：重构 layers 的 ModuleList
pruned_layers = torch.nn.ModuleList(
    [merged_model.model.layers[i] for i in keep_indices]
)
merged_model.model.layers = pruned_layers

# 2. 【核心修复】同步更新 config 中的所有层数及 layer_types 列表
configs_to_update = [merged_model.config]
if (
    hasattr(merged_model.config, "text_config")
    and merged_model.config.text_config is not None
):
    configs_to_update.append(merged_model.config.text_config)

for cfg in configs_to_update:
    cfg.num_hidden_layers = len(keep_indices)

    # 切割 layer_types 列表以匹配新的层数
    if hasattr(cfg, "layer_types") and cfg.layer_types is not None:
        cfg.layer_types = [cfg.layer_types[i] for i in keep_indices]

print(f"剪枝完成！新模型总层数: {len(merged_model.model.layers)} 层")

print(f"💾 阶段 4/4: 正在保存硬剪枝后的全新模型至 {OUTPUT_PRUNED_PATH}...")
os.makedirs(OUTPUT_PRUNED_PATH, exist_ok=True)
merged_model.save_pretrained(OUTPUT_PRUNED_PATH, safe_serialization=True)
tokenizer.save_pretrained(OUTPUT_PRUNED_PATH)

print("🎉 恭喜！硬剪枝独立模型已成功保存导出！")