File size: 6,671 Bytes
6eb3ab3 1b96d67 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | ---
license: apache-2.0
---
```python
import json
import sys
from pathlib import Path
import torch
from transformers import (
AutoConfig,
AutoProcessor,
GenerationConfig,
)
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForConditionalGeneration
SEED = 42
MODEL_ID = "openbmb/MiniCPM-OCR-Jev"
OUTPUT_DIR = Path(__file__).resolve().parent / "tiny-random-minicpm-ocr-jev"
# Marker written into the cached config so a stale/incompatible cache is rebuilt
# instead of silently reused.
CACHE_MARKER = "tiny-minicpm-ocr-jev-v1"
MAX_PARAMS = 40_000_000
MAX_BYTES = 180 * 1024 * 1024
def build_tiny_config():
"""Load the real config and shrink every scale parameter in place.
The full ``vocab_size`` (248320) is preserved so the real tokenizer and the
vision/image special-token ids (image/video/vision-start/vision-end) stay
valid for the image-text-to-text path.
"""
cfg = AutoConfig.from_pretrained(MODEL_ID)
tc = cfg.text_config
tc.hidden_size = 64
tc.intermediate_size = 128
tc.num_hidden_layers = 4
tc.num_attention_heads = 4
tc.num_key_value_heads = 2
tc.head_dim = 32
# Preserve the real linear/full attention mix (gated-delta-net + attention),
# with one full_attention layer like the full_attention_interval=4 pattern.
tc.layer_types = ["linear_attention", "linear_attention", "linear_attention", "full_attention"]
tc.linear_num_key_heads = 2
tc.linear_num_value_heads = 4
tc.linear_key_head_dim = 16
tc.linear_value_head_dim = 16
# keep tc.linear_conv_kernel_dim (4)
tc.max_position_embeddings = 4096
# keep tc.use_cache (original ships False; generation still works at this scale)
# partial_rotary_factor 0.25 * head_dim(32) -> rotary dim 8 -> freqs len 4;
# interleaved mrope_section must sum to rotary_dim // 2 == 4.
if isinstance(tc.rope_parameters, dict):
tc.rope_parameters["mrope_section"] = [2, 1, 1]
tc.rope_parameters["partial_rotary_factor"] = 0.25
vc = cfg.vision_config
vc.depth = 2
vc.hidden_size = 64
vc.intermediate_size = 128
vc.num_heads = 4
# out_hidden_size must match the text hidden_size (image embeds merge into
# the decoder embedding space).
vc.out_hidden_size = tc.hidden_size
# num_position_embeddings must be a perfect square (num_grid_per_side =
# int(sqrt(num_position_embeddings))); 256 -> 16 per side.
vc.num_position_embeddings = 256
# keep patch_size(16), spatial_merge_size(2), temporal_patch_size(2), in_channels(3)
# Force float32 at every effective precision level for a robust CPU forward.
cfg.dtype = "float32"
tc.dtype = "float32"
vc.dtype = "float32"
return cfg
def _count_and_measure(model):
n_params = 0
n_bytes = 0
for p in model.parameters():
n_params += p.numel()
n_bytes += p.numel() * p.element_size()
for b in model.buffers():
n_bytes += b.numel() * b.element_size()
return n_params, n_bytes
def _cache_is_valid():
cfg_path = OUTPUT_DIR / "config.json"
weights = (OUTPUT_DIR / "model.safetensors").exists() or (OUTPUT_DIR / "pytorch_model.bin").exists()
if not cfg_path.exists() or not weights:
return False
try:
data = json.loads(cfg_path.read_text())
except Exception:
return False
if data.get("_tiny_cache_marker") != CACHE_MARKER:
return False
if data.get("model_type") != "qwen3_5":
return False
if data.get("architectures") != ["Qwen3_5ForConditionalGeneration"]:
return False
return True
def create_tiny_model():
if _cache_is_valid():
print(f"Reusing valid cache at {OUTPUT_DIR}")
return OUTPUT_DIR
torch.manual_seed(SEED)
config = build_tiny_config()
# Instantiate the real architecture class with random weights.
with torch.device("cpu"):
model = Qwen3_5ForConditionalGeneration(config)
model = model.to(torch.float32)
model.eval()
n_params, n_bytes = _count_and_measure(model)
print(f"params={n_params:,} bytes={n_bytes / 1024 / 1024:.1f}MB")
if n_params > MAX_PARAMS:
raise RuntimeError(f"Parameter budget exceeded: {n_params:,} > {MAX_PARAMS:,}")
if n_bytes > MAX_BYTES:
raise RuntimeError(f"Memory budget exceeded: {n_bytes:,} > {MAX_BYTES:,}")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
model.save_pretrained(OUTPUT_DIR, safe_serialization=True)
# Save processor assets (tokenizer, image processor, chat template).
processor = AutoProcessor.from_pretrained(MODEL_ID)
processor.save_pretrained(OUTPUT_DIR)
# Save a usable generation config.
try:
gen = GenerationConfig.from_pretrained(MODEL_ID)
except Exception:
gen = GenerationConfig()
gen.save_pretrained(OUTPUT_DIR)
# Stamp the cache marker and assert float32 precision fields survived.
cfg_path = OUTPUT_DIR / "config.json"
data = json.loads(cfg_path.read_text())
data["_tiny_cache_marker"] = CACHE_MARKER
data["dtype"] = "float32"
if isinstance(data.get("text_config"), dict):
data["text_config"]["dtype"] = "float32"
if isinstance(data.get("vision_config"), dict):
data["vision_config"]["dtype"] = "float32"
# ``Qwen3_5VisionConfig.model_type`` class attribute is "qwen3_5", so
# ``to_dict()`` drops the real "qwen3_5_vision" marker; restore it.
data["vision_config"]["model_type"] = "qwen3_5_vision"
cfg_path.write_text(json.dumps(data, indent=2))
print(f"Saved tiny model to {OUTPUT_DIR}")
return OUTPUT_DIR
def verify_identity():
"""Assert the saved config preserves the real architecture identity."""
orig = AutoConfig.from_pretrained(MODEL_ID)
tiny = AutoConfig.from_pretrained(OUTPUT_DIR)
assert tiny.model_type == orig.model_type == "qwen3_5", tiny.model_type
assert tiny.architectures == orig.architectures == ["Qwen3_5ForConditionalGeneration"]
assert tiny.text_config.model_type == orig.text_config.model_type == "qwen3_5_text"
assert tiny.vision_config.model_type == orig.vision_config.model_type == "qwen3_5_vision"
assert tiny.image_token_id == orig.image_token_id
assert tiny.vision_start_token_id == orig.vision_start_token_id
assert tiny.text_config.vocab_size == orig.text_config.vocab_size
# linear + full attention mix preserved
assert "linear_attention" in tiny.text_config.layer_types
assert "full_attention" in tiny.text_config.layer_types
print("Architecture identity OK")
if __name__ == "__main__":
out = create_tiny_model()
verify_identity()
print(out)
``` |