File size: 2,756 Bytes
0695d0b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
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 左右。")
|