optimum-intel-internal-testing/tiny-random-minicpm-ocr-jev

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)