Gemma4 kerashub to safetensors checkpoint conversion

[kharshith-k]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces Hugging Face export utilities, unit tests, and verification scripts for the Gemma4 model in KerasHub, supporting both text-only and multimodal (vision and audio) configurations. The feedback highlights several key improvement opportunities: updating .gamma to .scale on normalization layers for Keras 3 compatibility, accessing configuration attributes directly from the vision and audio encoders rather than traversing internal layers, removing unused imports in the test file, and initializing st_keys in the verification script to prevent a potential NameError during exception handling.

[gemini-code-assist]

On line 483, conv_block.norm.gamma is used to retrieve the scale parameter of the LayerNormalization layer. However, in Keras 3, LayerNormalization uses scale as the attribute name for its scale weights (and bias for its center weights). Furthermore, the rest of this file consistently uses .scale to access normalization weights (e.g., block.pre_attention_norm.scale, backbone.get_layer("final_normalization").scale). Using .gamma here is inconsistent and may fail in Keras 3. It should be updated to .scale.

Suggested change

	weights_dict[f"{hf_conv_pfx}.norm.weight"] = conv_block.norm.gamma
	weights_dict[f"{hf_conv_pfx}.norm.weight"] = conv_block.norm.scale

[gemini-code-assist]

In _build_vision_config, the code retrieves the internal image_encoder layer and accesses its private attributes (image_encoder.patch_embedder.position_embedding_table.shape[1]) to get the position embedding size. However, Gemma4VisionEncoder already exposes position_embedding_size as a direct attribute (self.position_embedding_size). Accessing it directly is cleaner, safer, and avoids deep attribute traversal.

def _build_vision_config(vision_encoder):
    """Build the ``vision_config`` dict from a Gemma4VisionEncoder."""
    return {
        "model_type": "gemma4_vision_model",
        "num_hidden_layers": vision_encoder.num_layers,
        "num_attention_heads": vision_encoder.num_heads,
        "num_key_value_heads": vision_encoder.num_key_value_heads,
        "hidden_size": vision_encoder.hidden_dim,
        "intermediate_size": vision_encoder.intermediate_dim,
        "head_dim": vision_encoder.head_dim,
        "patch_size": vision_encoder.patch_size,
        "pooling_kernel_size": vision_encoder.pool_size,
        "position_embedding_size": vision_encoder.position_embedding_size,
        "rms_norm_eps": vision_encoder.layer_norm_epsilon,
        "use_clipped_linears": vision_encoder.use_clipped_linears,
        "standardize": vision_encoder.standardize,
        "rope_parameters": {
            "rope_theta": vision_encoder.rope_max_wavelength,
        },
    }

[gemini-code-assist]

In _build_audio_config, the code accesses private/internal layers of audio_encoder (such as audio_encoder.conformer_blocks[0].attention.attn and audio_encoder.subsample_conv_projection) to retrieve configuration parameters. However, Gemma4AudioEncoder already exposes almost all of these properties directly as attributes on the class instance (e.g., chunk_size, context_left, context_right, logit_cap, invalid_logit_value, conv_kernel_size, residual_weight, gradient_clipping, sscp_conv_channels, norm_eps). Accessing these attributes directly is much cleaner, more robust, and respects the encapsulation of the class.

def _build_audio_config(audio_encoder):
    """Build the ``audio_config`` dict from a Gemma4AudioEncoder."""
    return {
        "model_type": "gemma4_audio_model",
        "hidden_size": audio_encoder.hidden_size,
        "num_attention_heads": audio_encoder.num_heads,
        "num_hidden_layers": audio_encoder.num_layers,
        "attention_chunk_size": audio_encoder.chunk_size,
        "attention_context_left": audio_encoder.context_left,
        "attention_context_right": audio_encoder.context_right,
        "attention_logit_cap": audio_encoder.logit_cap,
        "attention_invalid_logits_value": audio_encoder.invalid_logit_value,
        "conv_kernel_size": audio_encoder.conv_kernel_size,
        "residual_weight": audio_encoder.residual_weight,
        "gradient_clipping": audio_encoder.gradient_clipping,
        "subsampling_conv_channels": list(audio_encoder.sscp_conv_channels),
        "output_proj_dims": audio_encoder.output_proj_dims,
        "rms_norm_eps": audio_encoder.norm_eps,
    }

[gemini-code-assist]

The imports AutoConfig, AutoModel, and AutoTokenizer from transformers are imported but never used anywhere in gemma4_test.py. Only AutoModelForCausalLM is used. Removing these unused imports will clean up the code and improve maintainability.

Suggested change

	from transformers import AutoConfig
	from transformers import AutoModel
	from transformers import AutoModelForCausalLM
	from transformers import AutoTokenizer
	from transformers import AutoModelForCausalLM

[gemini-code-assist]

In run_structural_smoke_test, the variable st_keys is defined inside the try block on line 471. If an exception occurs during the safetensors inspection (e.g., if safetensors is not installed or safe_open fails), the except block on line 481 will execute, but the script will continue to the next try block on line 485. In that block, missing_from_hf = [k for k in st_keys if k not in hf_sd_keys] is executed, which will raise a NameError: name 'st_keys' is not defined because st_keys was never initialized. To prevent this secondary crash, initialize st_keys = [] before the first try block.

Suggested change

	try:
	from safetensors import safe_open
	st_total = 0
	with safe_open(st_path, framework="pt") as f:
	st_keys = sorted(f.keys())
	for k in st_keys:
	st_total += f.get_tensor(k).numel()
	keras_params = model.count_params()
	count_match = st_total == keras_params
	print(
	f" {'✓' if count_match else '✗'} Parameter count: "
	f"Keras={keras_params:,}, safetensors={st_total:,}"
	)
	all_ok = all_ok and count_match
	except Exception as st_exc:
	print(f" ✗ Safetensors inspection failed: {st_exc}")
	all_ok = False
	st_keys = []
	try:
	from safetensors import safe_open
	st_total = 0
	with safe_open(st_path, framework="pt") as f:
	st_keys = sorted(f.keys())
	for k in st_keys:
	st_total += f.get_tensor(k).numel()
	keras_params = model.count_params()
	count_match = st_total == keras_params
	print(
	f" ✓ Parameter count: "
	f"Keras={keras_params:,}, safetensors={st_total:,}"
	)
	all_ok = all_ok and count_match
	except Exception as st_exc:
	print(f" ✗ Safetensors inspection failed: {st_exc}")
	all_ok = False