import os
import shutil
import torch
from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig

# === 1. 路径设置 ===
INPUT_PATH = "/home/nashen/deepseek-ocr/DeepSeek-OCR-master/DeepSeek-OCR-vllm/model/"
OUTPUT_PATH = "./DeepSeek-OCR-4bit-Quantized"  # 打包好的成品将保存在这里

# === 2. 黄金量化配置 (直接复用你测试成功的配置) ===
print("⏳ 正在加载并量化模型 (这可能需要几分钟)...")

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    
    # 🛡️ 核心白名单：确保保存后的模型也有这层保护
    llm_int8_skip_modules=[
        "sam_model", "model.sam_model",
        "vision_model", "model.vision_model",
        "projector", "model.projector",
        "lm_head", "embed_tokens"
    ]
)

# === 3. 加载模型 ===
model = AutoModel.from_pretrained(
    INPUT_PATH,
    trust_remote_code=True,
    quantization_config=quantization_config,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(INPUT_PATH, trust_remote_code=True)

print("✅ 模型加载完成，准备保存...")

# === 4. 保存模型 (核心步骤) ===
# 这会将 4-bit 权重和 quantization_config 写入 config.json
model.save_pretrained(OUTPUT_PATH, safe_serialization=True)
tokenizer.save_pretrained(OUTPUT_PATH)

print(f"✅ 权重已保存至: {OUTPUT_PATH}")

# === 5. 关键：复制自定义代码文件 ===
# 因为 DeepSeek-OCR 是 "Remote Code" 模型，必须手动把 python 代码拷过去
# 否则用户加载时会报错 "class not found"

print("📦 正在复制 Python 架构文件...")
files_to_copy = [
    "configuration_deepseek_v2.py",
    "modeling_deepseekv2.py",
    "modeling_deepseekocr.py",
    "deepencoder.py",
    # 如果有其他 .py 文件或 .json (如 processor_config.json) 也要拷过去
    "processor_config.json", 
    "config.json", # 覆盖一下以防万一，但 save_pretrained 通常会处理
    "special_tokens_map.json",
    "tokenizer_config.json",
    "tokenizer.json" 
]

# 自动扫描并复制 input 目录下的所有 .py 文件 (更安全)
for filename in os.listdir(INPUT_PATH):
    if filename.endswith(".py") or filename.endswith(".json"):
        src = os.path.join(INPUT_PATH, filename)
        dst = os.path.join(OUTPUT_PATH, filename)
        
        # config.json 和 tokenizer 文件可能已经被 save_pretrained 生成了，跳过覆盖
        if os.path.exists(dst) and "config" in filename:
            continue
            
        shutil.copy2(src, dst)

print(f"🎉 打包完成！成品位于: {OUTPUT_PATH}")
print(f"   该文件夹现在的体积应该是 ~2.7 GB 左右。")