| --- |
| license: apache-2.0 |
| --- |
| |
| ```python |
| |
| import argparse |
| from pathlib import Path |
| |
| |
| DEFAULT_SOURCE_MODEL = "google/gemma-4-E2B-it" |
| |
| |
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Generate and optionally export a configurable tiny-random Gemma4 model.", |
| ) |
| parser.add_argument("--output-dir", type=Path, required=True) |
| parser.add_argument("--source-model", default=DEFAULT_SOURCE_MODEL) |
| parser.add_argument("--hidden-size", type=int, default=32) |
| parser.add_argument("--head-dim", type=int, default=8) |
| parser.add_argument("--num-attention-heads", type=int, default=4) |
| parser.add_argument("--intermediate-size", type=int, default=64) |
| parser.add_argument("--text-layers", type=int, default=3) |
| parser.add_argument("--audio-layers", type=int, default=1) |
| parser.add_argument("--vision-layers", type=int, default=1) |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument("--export-dir", type=Path) |
| parser.add_argument("--smoke-test", action="store_true") |
| parser.add_argument("--device", default="CPU") |
| parser.add_argument("--attention-backend", choices=("PA", "SDPA"), default="SDPA") |
| return parser.parse_args() |
| |
| |
| def validate_args(args: argparse.Namespace) -> None: |
| if args.hidden_size != args.head_dim * args.num_attention_heads: |
| raise ValueError( |
| "hidden-size must equal head-dim * num-attention-heads: " |
| f"{args.hidden_size} != {args.head_dim} * {args.num_attention_heads}" |
| ) |
| if args.num_attention_heads < 2: |
| raise ValueError("num-attention-heads must be at least 2") |
| if args.text_layers != 3: |
| raise ValueError("text-layers must remain 3 for the model-card layer pattern") |
| if args.intermediate_size < args.hidden_size: |
| raise ValueError("intermediate-size must be at least hidden-size") |
| if args.smoke_test and args.export_dir is None: |
| raise ValueError("--smoke-test requires --export-dir") |
| |
| |
| def generate_model(args: argparse.Namespace) -> None: |
| import torch |
| from transformers import AutoProcessor, Gemma4Config, Gemma4ForConditionalGeneration |
| |
| torch.manual_seed(args.seed) |
| config = Gemma4Config.from_pretrained(args.source_model) |
| |
| config.audio_config.hidden_size = args.hidden_size |
| config.audio_config.num_attention_heads = args.num_attention_heads |
| config.audio_config.num_hidden_layers = args.audio_layers |
| config.audio_config.output_proj_dims = args.hidden_size |
| config.audio_config.dtype = "float32" |
| |
| config.text_config.global_head_dim = args.head_dim |
| config.text_config.head_dim = args.head_dim |
| config.text_config.hidden_size = args.hidden_size |
| config.text_config.hidden_size_per_layer_input = 1 |
| config.text_config.intermediate_size = args.intermediate_size |
| config.text_config.num_attention_heads = args.num_attention_heads |
| config.text_config.num_key_value_heads = max(1, args.num_attention_heads // 2) |
| config.text_config.num_hidden_layers = args.text_layers |
| config.text_config.layer_types = ["sliding_attention", "full_attention", "full_attention"] |
| config.text_config.num_kv_shared_layers = 1 |
| config.text_config.dtype = "float32" |
| |
| config.vision_config.default_output_length = 70 |
| config.vision_config.head_dim = args.head_dim |
| config.vision_config.hidden_size = args.hidden_size |
| config.vision_config.intermediate_size = args.intermediate_size |
| config.vision_config.num_attention_heads = args.num_attention_heads |
| config.vision_config.num_hidden_layers = args.vision_layers |
| config.vision_config.num_key_value_heads = args.num_attention_heads |
| config.vision_config.patch_size = 16 |
| config.vision_config.dtype = "float32" |
| |
| model = Gemma4ForConditionalGeneration(config) |
| model.eval() |
| |
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| model.save_pretrained(args.output_dir) |
| processor = AutoProcessor.from_pretrained(args.source_model, padding_side="left", truncation_side="left") |
| processor.save_pretrained(args.output_dir) |
| |
| parameter_count = sum(parameter.numel() for parameter in model.parameters()) |
| print(f"Saved {parameter_count:,}-parameter model to {args.output_dir}") |
| |
| from transformers import AutoProcessor, Gemma4ForConditionalGeneration |
| |
| messages = [ |
| { |
| "role": "user", "content": [ |
| {"type": "image", |
| "url": "https://raw.githubusercontent.com/google-gemma/cookbook/refs/heads/main/apps/sample-data/GoldenGate.png"}, |
| {"type": "text", "text": "What is shown in this image?"} |
| ] |
| } |
| ] |
| |
| processor = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") |
| model = Gemma4ForConditionalGeneration.from_pretrained( |
| args.output_dir, |
| dtype="auto", |
| device_map="auto" |
| ) |
| |
| # Process input |
| inputs = processor.apply_chat_template( |
| messages, |
| tokenize=True, |
| return_dict=True, |
| return_tensors="pt", |
| add_generation_prompt=True, |
| ).to(model.device) |
| input_len = inputs["input_ids"].shape[-1] |
| |
| # Generate output |
| outputs = model.generate(**inputs, max_new_tokens=512) |
| print("VLM infer OK") |
| |
| |
| def export_model(model_dir: Path, export_dir: Path) -> None: |
| import openvino |
| import openvino_tokenizers |
| from optimum.intel.openvino import OVModelForVisualCausalLM |
| from transformers import AutoProcessor, Gemma4ForConditionalGeneration |
| |
| |
| |
| processor = AutoProcessor.from_pretrained(model_dir, padding_side="left", truncation_side="left") |
| ov_model = OVModelForVisualCausalLM.from_pretrained( |
| model_dir, |
| compile=False, |
| device="CPU", |
| export=True, |
| load_in_8bit=False, |
| ) |
| |
| processor.image_processor.size ={ |
| "height": 32, |
| "width": 32 |
| } |
| |
| export_dir.mkdir(parents=True, exist_ok=True) |
| ov_model.save_pretrained(export_dir) |
| processor.save_pretrained(export_dir) |
| ov_tokenizer, ov_detokenizer = openvino_tokenizers.convert_tokenizer( |
| processor.tokenizer, |
| with_detokenizer=True, |
| ) |
| openvino.save_model(ov_tokenizer, export_dir / "openvino_tokenizer.xml") |
| openvino.save_model(ov_detokenizer, export_dir / "openvino_detokenizer.xml") |
| print(f"Exported OpenVINO model to {export_dir}") |
| |
| |
| def smoke_test(export_dir: Path, device: str, attention_backend: str) -> None: |
| import numpy as np |
| import openvino |
| from openvino_genai import VLMPipeline |
| from transformers import Gemma4ForConditionalGeneration |
| |
| |
| pipeline = VLMPipeline(export_dir, device, ATTENTION_BACKEND=attention_backend) |
| text_result = pipeline.generate("Hello", max_new_tokens=1, do_sample=False) |
| print(f"Text smoke test passed: {text_result.texts!r}") |
| |
| sampling_rate = 16_000 |
| timestamps = np.arange(sampling_rate, dtype=np.float32) / sampling_rate |
| audio = 0.5 * np.sin(2 * np.pi * 440 * timestamps) + 0.25 * np.sin(2 * np.pi * 880 * timestamps) |
| audio_result = pipeline.generate( |
| "Describe this audio.<|audio|>", |
| audios=[openvino.Tensor(audio.astype(np.float32))], |
| max_new_tokens=1, |
| do_sample=False, |
| ) |
| print(f"Audio smoke test passed: {audio_result.texts!r}") |
| |
| |
| def main() -> None: |
| args = parse_args() |
| validate_args(args) |
| generate_model(args) |
| if args.export_dir is not None: |
| export_model(args.output_dir, args.export_dir) |
| if args.smoke_test: |
| smoke_test(args.export_dir, args.device, args.attention_backend) |
| |
| |
| if __name__ == "__main__": |
| main() |
| |
| ``` |
|
|
