optimum-intel-internal-testing/tiny-random-ministral3

"""Create a tiny random `ministral3` (Mistral3 VLM) fixture

Produces a Mistral3ForConditionalGeneration wrapper whose text_config.model_type is
"ministral3" (YaRN rope) around a pixtral vision tower.

Requirements: transformers >= 5.0.0 (ministral3 model_type), torch, pillow.

Usage:
    python create_tiny_ministral3_model.py
    # writes ./tiny-random-ministral3, then upload it:
    huggingface-cli upload optimum-intel-internal-testing/tiny-random-ministral3 tiny-random-ministral3 .
"""

import sys
from pathlib import Path

import torch
from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor


SEED = 42
MODEL_ID = "mistralai/Ministral-3-3B-Reasoning-2512"
OUTPUT_DIR = Path(__file__).resolve().parent / "tiny-random-ministral3"

MAX_PARAMS = 25_000_000
MAX_BYTES = 100 * 1024 * 1024


def build_config():
    config = AutoConfig.from_pretrained(MODEL_ID)

    config.tie_word_embeddings = False
    config.text_config.tie_word_embeddings = False

    config.text_config.num_hidden_layers = 2
    config.text_config.hidden_size = 64
    config.text_config.intermediate_size = 128
    config.text_config.num_attention_heads = 4
    config.text_config.num_key_value_heads = 2
    config.text_config.head_dim = 16
    config.text_config.max_position_embeddings = 512

    # Preserve the YaRN rope path: keep rope_type == "yarn" and all yarn fields, only rescale
    # original_max_position_embeddings so factor * original == max_position_embeddings.
    rope_scaling = dict(config.text_config.rope_scaling)
    assert rope_scaling["rope_type"] == "yarn", rope_scaling
    rope_scaling["original_max_position_embeddings"] = 32
    config.text_config.rope_scaling = rope_scaling

    config.vision_config.num_hidden_layers = 2
    config.vision_config.hidden_size = 64
    config.vision_config.intermediate_size = 128
    config.vision_config.num_attention_heads = 4
    config.vision_config.head_dim = 16
    config.vision_config.image_size = 56

    for subconfig in (config, config.text_config, config.vision_config):
        subconfig.dtype = "float32"
        subconfig.torch_dtype = "float32"

    return config


def main():
    torch.manual_seed(SEED)

    config = build_config()
    model = AutoModelForImageTextToText.from_config(config).float().eval()

    num_params = sum(p.numel() for p in model.parameters())
    num_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    print(f"params={num_params:,}  bytes={num_bytes:,} ({num_bytes / 1024 / 1024:.2f} MB)")
    if num_params > MAX_PARAMS:
        sys.exit(f"param budget exceeded: {num_params:,} > {MAX_PARAMS:,}")
    if num_bytes > MAX_BYTES:
        sys.exit(f"memory budget exceeded: {num_bytes:,} > {MAX_BYTES:,}")

    processor = AutoProcessor.from_pretrained(MODEL_ID)
    processor.image_processor.size = {"longest_edge": 56}

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(OUTPUT_DIR, safe_serialization=True)
    processor.save_pretrained(OUTPUT_DIR)

    # Validate architecture identity + task path on the reloaded artifact.
    reloaded_config = AutoConfig.from_pretrained(OUTPUT_DIR)
    assert reloaded_config.model_type == "mistral3", reloaded_config.model_type
    assert reloaded_config.text_config.model_type == "ministral3", reloaded_config.text_config.model_type
    assert reloaded_config.text_config.rope_scaling["rope_type"] == "yarn"

    reloaded = AutoModelForImageTextToText.from_pretrained(OUTPUT_DIR).eval()
    reloaded_processor = AutoProcessor.from_pretrained(OUTPUT_DIR)

    from PIL import Image

    image = Image.new("RGB", (56, 56), color=(127, 127, 127))
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": "Describe the image."},
            ],
        }
    ]
    prompt = reloaded_processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
    inputs = reloaded_processor(text=prompt, images=[image], return_tensors="pt")
    with torch.no_grad():
        generated = reloaded.generate(**inputs, max_new_tokens=5, do_sample=False)
    print("generate ok, output shape:", tuple(generated.shape))

    print(f"\nSaved tiny ministral3 model to: {OUTPUT_DIR}")


if __name__ == "__main__":
    main()