import os
import torch
torch.set_default_dtype(torch.float32)
from transformers import AutoConfig, AutoProcessor
from transformers.models.deepseek_ocr2 import DeepseekOcr2ForConditionalGeneration
# === Step 1: Load and shrink the original config ===
MODEL_ID = "deepseek-community/DeepSeek-OCR-2"
OUTPUT_DIR = "./tiny-deepseek-ocr-2"
print(f"Loading config from {MODEL_ID}...")
config = AutoConfig.from_pretrained(MODEL_ID)
# Tiny DeepseekV2 (MoE) text decoder (reduced from hidden_size=1280, 12 layers, 10 heads).
text_config = config.text_config
text_config.hidden_size = 128 # from 1280
text_config.intermediate_size = 256 # from 6848 (dense MLP)
text_config.moe_intermediate_size = 128 # from 896 (per-expert MLP)
text_config.num_hidden_layers = 2 # from 12
text_config.num_attention_heads = 4 # from 10
text_config.num_key_value_heads = 4 # from 10
text_config.head_dim = 32 # from 128
text_config.n_routed_experts = 4 # from 64
text_config.num_experts_per_tok = 2 # from 6
text_config.n_shared_experts = 1 # from 2
# First layer stays dense, the rest are MoE ("sparse"); length must equal num_hidden_layers.
text_config.mlp_layer_types = ["dense", "sparse"]
# Tiny CLIP-style vision encoder (reduced from hidden_size=896, 24 layers, 14 heads).
encoder_config = config.vision_config.encoder_config
encoder_config.hidden_size = 128 # from 896 (must match sam downsample_channels[1] below)
encoder_config.intermediate_size = 256 # from 4864
encoder_config.num_hidden_layers = 2 # from 24
encoder_config.num_attention_heads = 4 # from 14
encoder_config.num_key_value_heads = 2 # from 2
encoder_config.layer_types = ["full_attention"] * encoder_config.num_hidden_layers
# Tiny SAM ViT vision encoder (reduced from hidden_size=768, 12 layers, 12 heads).
sam_config = config.vision_config.sam_config
sam_config.hidden_size = 64 # from 768
sam_config.num_hidden_layers = 2 # from 12
sam_config.num_attention_heads = 4 # from 12
sam_config.mlp_dim = 128 # from 3072
sam_config.output_channels = 64 # from 256
sam_config.downsample_channels = [64, 128] # from [512, 896]; last must == encoder hidden_size
sam_config.global_attn_indexes = [] # windowed attention only, to keep compute small
# image_size (1024), patch_size (16) and window_size (14) are kept so the preprocessing
# (global 1024px view + 768px crop tiles) stays compatible with the real model.
# === Step 2: Create model from config ===
print("Creating tiny DeepSeek-OCR-2 model...")
model = DeepseekOcr2ForConditionalGeneration(config)
model.eval()
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total_params:,} ({total_params * 4 / 1024 / 1024:.2f} MB in float32)")
# === Step 3: Load processor (image processor + tokenizer) from the original model ===
print(f"Loading processor from {MODEL_ID}...")
processor = AutoProcessor.from_pretrained(MODEL_ID)
# === Step 4: Save model and processor ===
os.makedirs(OUTPUT_DIR, exist_ok=True)
print(f"Saving tiny model to {OUTPUT_DIR}...")
model.save_pretrained(OUTPUT_DIR, safe_serialization=False)
processor.save_pretrained(OUTPUT_DIR)
print(f"Done! Tiny DeepSeek-OCR-2 model saved to {OUTPUT_DIR}")
