demo/finalize_save.py

# ====================================================================================
# finalize_save.py (v1.1 - Memory Optimized)
# ====================================================================================

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gc # 导入垃圾回收模块

print("--- 运行最终模型保存脚本 (内存优化版) ---")

# --- [STEP 1] 定义模型和路径 ---
MODEL_ID = "./gemma-3-4b-it-qat-q4_0-unquantized"
NEW_MODEL_ID = "./gemma-3-4b-it-qat-q4_0-unquantized-precise-abliterated" 
WEIGHTS_FILE = "./consolidated_precise_abliterated_weights.pth"

print(f"原始模型 ID: {MODEL_ID}")
print(f"权重文件路径: {WEIGHTS_FILE}")
print(f"最终保存路径: {NEW_MODEL_ID}")


# --- [STEP 2] 加载一个干净的原始模型到 CPU ---
print("\n[STEP 2] 正在 CPU 上加载一个干净的原始模型架构...")
clean_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    local_files_only=True,
    torch_dtype=torch.bfloat16,
    device_map="cpu" 
)
print("[SUCCESS] 干净的模型加载完成。")


# --- [STEP 3] 加载之前保存的状态字典 ---
print(f"\n[STEP 3] 正在从 {WEIGHTS_FILE} 加载手术后的权重...")
consolidated_weights = torch.load(WEIGHTS_FILE, map_location="cpu")
print("[SUCCESS] 权重文件加载完成。")


# --- [STEP 4] 将权重应用到干净的模型上 ---
print("\n[STEP 4] 正在将权重应用到模型...")
clean_model.load_state_dict(consolidated_weights)
print("[SUCCESS] 权重成功应用。")


# --- [新增步骤] 释放内存 ---
print("\n[内存优化] 正在释放已加载的权重字典以节省内存...")
del consolidated_weights # 删除不再需要的权重字典
gc.collect() # 强制执行垃圾回收
print("[SUCCESS] 内存已释放。")


# --- [STEP 5] 保存最终模型和分词器 ---
print(f"\n[STEP 5] 正在将最终的、完整的模型保存到: {NEW_MODEL_ID}")
clean_model.save_pretrained(NEW_MODEL_ID)

# 同时，也需要保存分词器
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
tokenizer.save_pretrained(NEW_MODEL_ID)

print(f"\n[全部完成!] 您的手术后模型已成功保存到: {NEW_MODEL_ID}")