from transformers import AutoModel
import torch

# 你的模型路径
MODEL_PATH = "/home/nashen/deepseek-ocr/DeepSeek-OCR-master/DeepSeek-OCR-vllm/model/"

print(f"⏳ 正在加载模型结构 (不加载权重，速度很快)...")

# 技巧：使用 meta device 或只加载 config，避免消耗显存
try:
    # 尝试加载到 CPU，不进行量化，只为了看名字
    model = AutoModel.from_pretrained(
        MODEL_PATH, 
        trust_remote_code=True, 
        device_map="cpu", # 强制 CPU，不占显存
        torch_dtype=torch.float16
    )
except Exception as e:
    print(f"❌ 加载出错: {e}")
    exit()

print(f"\n{'='*20} 模型层级结构 {'='*20}")
print("🔍 正在寻找可能的视觉层 (Vision) 和投影层 (Projector)...")
print(f"{'='*60}")

# 遍历所有层
vision_candidates = []
for name, module in model.named_modules():
    # 只打印顶级和二级模块，防止刷屏
    depth = name.count(".")
    if depth <= 2:
        # 获取类名，帮助判断它是干嘛的
        class_name = module.__class__.__name__
        print(f"🔹 层名: {name:<40} | 类型: {class_name}")
        
        # 智能识别可能的视觉层
        lower_name = name.lower()
        if "vision" in lower_name or "img" in lower_name or "encoder" in lower_name or "proj" in lower_name or "align" in lower_name:
            vision_candidates.append(name)

print(f"\n{'='*20} 🎯 建议加入白名单的层 {'='*20}")
print("请把以下名字加入 test_benchmark.py 的 llm_int8_skip_modules 列表中：")
for cand in vision_candidates:
    print(f'  "{cand}",')

print(f"{'='*60}")