Update README.md

README.md

@@ -3,68 +3,199 @@ license: apache-2.0
 ---
 
 ```python
-from transformers import (
-    AutoTokenizer,
-    Gemma4Config,
-    Gemma4ForConditionalGeneration,
-    Gemma4TextConfig,
-    Gemma4ForCausalLM,
-    Gemma4VisionConfig,
-    Gemma4AudioConfig,
-)
-
-
-def generate_vlm_model(output_dir="./tiny-random-gemma4"):
-    model_tr = Gemma4ForConditionalGeneration.from_pretrained("google/gemma-4-E2B-it")
-    config = model_tr.config
-
-    config.audio_config.hidden_size = 8
-    config.audio_config.num_attention_heads = 2
-    config.audio_config.num_hidden_layers = 1
-    config.audio_config.output_proj_dims = 8
-
-    config.text_config.global_head_dim = 4
-    config.text_config.head_dim = 4
-    config.text_config.hidden_size = 8
+
+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 = 32
-    config.text_config.num_attention_heads = 2
-    config.text_config.num_hidden_layers = 3
+    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 = 4
-    config.vision_config.hidden_size = 8
-    config.vision_config.intermediate_size = 32
-    config.vision_config.num_attention_heads = 2
-    config.vision_config.num_hidden_layers = 1
-    config.vision_config.num_key_value_heads = 2
-    config.vision_config.patch_size = 2
+    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()
 
-    model.save_pretrained(output_dir)
+    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
+
 
-    # Copy tokenizer from google/gemma-4-E2B-it
-    tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-E2B-it")
-    tokenizer.save_pretrained(output_dir)
+    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}")
 
-    # Estimate safetensors size
-    import os
+    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}")
 
-    safetensors_path = os.path.join(output_dir, "model.safetensors")
-    if os.path.exists(safetensors_path):
-        size_mb = os.path.getsize(safetensors_path) / (1024 * 1024)
-        print(f"  model.safetensors size: {size_mb:.1f} MB")
 
-    print(f"  VLM model saved to {output_dir}")
-    return model
+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__":
-    generate_vlm_model()
+    main()
 
 ```