| --- |
| license: apache-2.0 |
| --- |
| |
| ```python |
| """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() |
| ``` |
