| --- |
| 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) |
| |
| ``` |
