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