Views
No views yet
1
2import os
3
4import torch
5
6
7torch.set_default_dtype(torch.float32)
8
9from transformers import AutoConfig, AutoProcessor
10from transformers.models.deepseek_ocr2 import DeepseekOcr2ForConditionalGeneration
11
12
13# === Step 1: Load and shrink the original config ===
14MODEL_ID = "deepseek-community/DeepSeek-OCR-2"
15OUTPUT_DIR = "./tiny-deepseek-ocr-2"
16
17print(f"Loading config from {MODEL_ID}...")
18config = AutoConfig.from_pretrained(MODEL_ID)
19
20# Tiny DeepseekV2 (MoE) text decoder (reduced from hidden_size=1280, 12 layers, 10 heads).
21text_config = config.text_config
22text_config.hidden_size = 128 # from 1280
23text_config.intermediate_size = 256 # from 6848 (dense MLP)
24text_config.moe_intermediate_size = 128 # from 896 (per-expert MLP)
25text_config.num_hidden_layers = 2 # from 12
26text_config.num_attention_heads = 4 # from 10
27text_config.num_key_value_heads = 4 # from 10
28text_config.head_dim = 32 # from 128
29text_config.n_routed_experts = 4 # from 64
30text_config.num_experts_per_tok = 2 # from 6
31text_config.n_shared_experts = 1 # from 2
32# First layer stays dense, the rest are MoE ("sparse"); length must equal num_hidden_layers.
33text_config.mlp_layer_types = ["dense", "sparse"]
34
35# Tiny CLIP-style vision encoder (reduced from hidden_size=896, 24 layers, 14 heads).
36encoder_config = config.vision_config.encoder_config
37encoder_config.hidden_size = 128 # from 896 (must match sam downsample_channels[1] below)
38encoder_config.intermediate_size = 256 # from 4864
39encoder_config.num_hidden_layers = 2 # from 24
40encoder_config.num_attention_heads = 4 # from 14
41encoder_config.num_key_value_heads = 2 # from 2
42encoder_config.layer_types = ["full_attention"] * encoder_config.num_hidden_layers
43
44# Tiny SAM ViT vision encoder (reduced from hidden_size=768, 12 layers, 12 heads).
45sam_config = config.vision_config.sam_config
46sam_config.hidden_size = 64 # from 768
47sam_config.num_hidden_layers = 2 # from 12
48sam_config.num_attention_heads = 4 # from 12
49sam_config.mlp_dim = 128 # from 3072
50sam_config.output_channels = 64 # from 256
51sam_config.downsample_channels = [64, 128] # from [512, 896]; last must == encoder hidden_size
52sam_config.global_attn_indexes = [] # windowed attention only, to keep compute small
53# image_size (1024), patch_size (16) and window_size (14) are kept so the preprocessing
54# (global 1024px view + 768px crop tiles) stays compatible with the real model.
55
56# === Step 2: Create model from config ===
57print("Creating tiny DeepSeek-OCR-2 model...")
58model = DeepseekOcr2ForConditionalGeneration(config)
59model.eval()
60
61total_params = sum(p.numel() for p in model.parameters())
62print(f"Total parameters: {total_params:,} ({total_params * 4 / 1024 / 1024:.2f} MB in float32)")
63
64# === Step 3: Load processor (image processor + tokenizer) from the original model ===
65print(f"Loading processor from {MODEL_ID}...")
66processor = AutoProcessor.from_pretrained(MODEL_ID)
67
68# === Step 4: Save model and processor ===
69os.makedirs(OUTPUT_DIR, exist_ok=True)
70print(f"Saving tiny model to {OUTPUT_DIR}...")
71model.save_pretrained(OUTPUT_DIR, safe_serialization=False)
72processor.save_pretrained(OUTPUT_DIR)
73
74print(f"Done! Tiny DeepSeek-OCR-2 model saved to {OUTPUT_DIR}")