Add Gemma3n #2329
Add Gemma3n #2329
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @kanpuriyanawab, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request initiates the integration of the Gemma3n model into the keras_hub repository. It primarily focuses on establishing the core architectural components for both audio encoding and text decoding, which are crucial for the model's multi-modal capabilities. The changes provide the foundational Keras layers necessary to construct the full Gemma3n model, including its unique attention, feed-forward, and normalization mechanisms.
Highlights
- New Model Integration: This pull request introduces the foundational components for the Gemma3n model, laying the groundwork for a multi-modal Causal Language Model within the
keras_hubrepository. - Audio Encoder Architecture: A comprehensive audio encoder based on the Conformer architecture has been added. This includes specialized attention mechanisms with relative position embeddings, feed-forward networks, and lightweight convolutions, designed to process audio inputs effectively.
- Advanced Text Decoder Architecture: A sophisticated text decoder has been implemented, featuring novel components such as the 'Learned Augmented Residual Layer (Laurel)' and 'Alternating Updates (AltUp)' blocks. These, along with custom MLP and attention mechanisms, are designed to facilitate complex multimodal fusion and processing.
- Core Utility Components: A custom RMS Normalization layer (
Gemma3nRMSNorm) has been introduced, which serves as a fundamental building block used consistently across various Gemma3n model components for stable training.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces the Gemma3n model, including audio encoder, causal LM, layer normalization, and text decoder components. No specific style guide was provided, and no immediate issues requiring comments at the 'medium', 'high', or 'critical' severity levels were identified in the added code.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces the Gemma3n model, a significant and complex addition. While the implementation of the core components shows a good understanding of the architecture, there are several critical issues that must be addressed before this can be merged. The most severe issues include a NameError due to an undefined loop variable and multiple AttributeErrors from invalid keys in the get_config method within gemma3n_backbone.py. Additionally, the mobilenetv5_backbone.py file contains a large block of duplicated code that makes it invalid. Furthermore, this PR has major gaps in adhering to the repository's style guide: there are no tests for any of the new gemma3n components, and the main task model (Gemma3nForConditionalGeneration) is just a placeholder. These issues need to be resolved to ensure the model is correct, maintainable, and verifiable.
| for transformer_layer in self.transformer_layers: | ||
| start = i * hidden_size_per_layer_input |
There was a problem hiding this comment.
The loop variable i is not defined, which will raise a NameError when this code is executed. You should use enumerate to get both the index and the item from the iterable.
| for transformer_layer in self.transformer_layers: | |
| start = i * hidden_size_per_layer_input | |
| for i, transformer_layer in enumerate(self.transformer_layers): |
| "query_head_dim_normalize": self.query_head_dim_normalize, | ||
| "use_query_key_norm": self.use_query_key_norm, | ||
| "use_post_ffw_norm": self.use_post_ffw_norm, | ||
| "use_post_attention_norm": self.use_post_attention_norm, | ||
| "attention_logit_soft_cap": self.attention_logit_soft_cap, | ||
| "final_logit_soft_cap": self.final_logit_soft_cap, | ||
| "use_sliding_window_attention": ( | ||
| self.use_sliding_window_attention | ||
| ), | ||
| "sliding_window_size": self.sliding_window_size, | ||
| "local_rope_scaling_factor": self.local_rope_scaling_factor, | ||
| "global_rope_scaling_factor": self.global_rope_scaling_factor, |
There was a problem hiding this comment.
This get_config method includes several keys that do not correspond to attributes set on the Gemma3nBackbone instance (e.g., query_head_dim_normalize, use_sliding_window_attention, local_rope_scaling_factor). This will cause an AttributeError during serialization. These keys appear to be leftovers from another model and should be removed to ensure the configuration is correct.
| b_args["out_chs"] = adjust_channels(b_args["out_chs"]) | ||
|
|
||
| if block_type == "er": | ||
| block = EdgeResidualBlock( | ||
| in_chs=b_args["in_chs"], | ||
| out_chs=b_args["out_chs"], | ||
| exp_ratio=b_args["exp_ratio"], | ||
| kernel_size=int(b_args["dw_kernel_size_mid"]), | ||
| stride=b_args["stride"], | ||
| act_layer=b_args["act_layer"], | ||
| layer_norm_epsilon=layer_norm_epsilon, | ||
| name=block_name, | ||
| dtype=dtype, | ||
| ) | ||
| elif block_type == "uir": | ||
| block = UniversalInvertedResidualBlock( | ||
| in_chs=b_args["in_chs"], | ||
| out_chs=b_args["out_chs"], | ||
| exp_ratio=b_args["exp_ratio"], | ||
| dw_kernel_size_mid=b_args["dw_kernel_size_mid"], | ||
| dw_kernel_size_start=b_args.get( | ||
| "dw_kernel_size_start", 0 | ||
| ), | ||
| dw_kernel_size_end=b_args.get("dw_kernel_size_end", 0), | ||
| stride=b_args["stride"], | ||
| se_ratio=b_args.get("se_ratio"), | ||
| act_layer=b_args["act_layer"], | ||
| layer_norm_epsilon=layer_norm_epsilon, | ||
| layer_scale_init_value=layer_scale_init_value, | ||
| name=block_name, | ||
| dtype=dtype, | ||
| ) | ||
| elif block_type == "mqa": | ||
| # This block can also be made explicit for safety | ||
| b_args.pop("act_layer", None) | ||
| block = MobileAttentionBlock( | ||
| # in_chs=b_args["in_chs"], | ||
| out_chs=b_args["out_chs"], | ||
| stride=b_args["stride"], | ||
| num_heads=b_args["num_heads"], | ||
| key_dim=b_args["key_dim"], | ||
| value_dim=b_args["value_dim"], | ||
| dw_kernel_size=b_args["dw_kernel_size"], | ||
| kv_stride=b_args.get("kv_stride", 1), | ||
| layer_norm_epsilon=layer_norm_epsilon, | ||
| layer_scale_init_value=layer_scale_init_value, | ||
| name=block_name, | ||
| dtype=dtype, | ||
| ) | ||
| else: | ||
| raise ValueError(f"Unknown block type: {block_type}") | ||
|
|
||
| x = block(x) | ||
| current_in_chs = b_args["out_chs"] | ||
| feature_maps.append(x) | ||
|
|
||
| # Multi-Scale Fusion Adapter (MSFA) | ||
| msfa_input_maps = [feature_maps[i] for i in msfa_indices] | ||
|
|
||
| x = MobileNetV5MultiScaleFusionAdapter( | ||
| in_chs_list=[m.shape[channel_axis] for m in msfa_input_maps], | ||
| out_chs=num_features, | ||
| output_resolution=msfa_output_resolution, | ||
| layer_norm_epsilon=layer_norm_epsilon, | ||
| act_layer=act_layer, | ||
| name="msfa", | ||
| dtype=dtype, | ||
| )(msfa_input_maps) | ||
|
|
||
| super().__init__(inputs=inputs, outputs=x, dtype=dtype, **kwargs) | ||
|
|
||
| # Store config | ||
| self.block_args = block_args | ||
| self.stem_size = stem_size | ||
| self.stem_bias = stem_bias | ||
| self.msfa_indices = msfa_indices | ||
| self.msfa_output_resolution = msfa_output_resolution | ||
| self.num_features = num_features | ||
| self.image_shape = image_shape | ||
| self.act_layer = act_layer | ||
| self.layer_norm_epsilon = layer_norm_epsilon | ||
| self.layer_scale_init_value = layer_scale_init_value | ||
|
|
||
| def get_config(self): | ||
| config = super().get_config() | ||
| config.update( | ||
| { | ||
| "block_args": self.block_args, | ||
| "stem_size": self.stem_size, | ||
| "stem_bias": self.stem_bias, | ||
| "msfa_indices": self.msfa_indices, | ||
| "msfa_output_resolution": self.msfa_output_resolution, | ||
| "num_features": self.num_features, | ||
| "image_shape": self.image_shape, | ||
| "act_layer": self.act_layer, | ||
| "layer_norm_epsilon": self.layer_norm_epsilon, | ||
| "layer_scale_init_value": self.layer_scale_init_value, | ||
| } | ||
| ) | ||
| return config |
There was a problem hiding this comment.
The file from line 895 onwards contains duplicated definitions for _decode_block_str, decode_arch_def, and the MobileNetV5Backbone class. This makes the Python file invalid and will cause errors. Please remove the duplicated code.
| class Gemma3nForConditionalGeneration(CausalLM): | ||
| """An end-to-end multi-modal Gemma3n model for Causal LM.""" | ||
|
|
||
| backbone_cls = None # Would be Gemma3nBackbone | ||
| preprocessor_cls = Gemma3nCausalLMPreprocessor | ||
|
|
||
| def __init__(self, backbone, preprocessor=None, **kwargs): |
There was a problem hiding this comment.
This class is a placeholder and does not meet the requirements for a task model as defined in the style guide. The backbone_cls is not set, and the generate_step method is not implemented. A complete task model is required to integrate the backbone with a preprocessor and a task-specific head, providing a user-friendly API like generate() and fit(). 1
Please provide a full implementation for this class.
Style Guide References
Footnotes
| self.altup_num_inputs, use_bias=False, name="modality_router" | ||
| ) | ||
| self.router_norm = Gemma3nRMSNorm( | ||
| # self.hidden_size, |
There was a problem hiding this comment.
This commented-out code should be removed for clarity.
| # The specific fix is here: `scale` -> `with_scale` | ||
| self.v_norm = Gemma3nRMSNorm( | ||
| epsilon=self.rms_norm_eps, | ||
| with_scale=False, # Corrected keyword argument |
There was a problem hiding this comment.
This comment appears to be for the developer or reviewer and should be removed from the final code.
| with_scale=False, # Corrected keyword argument | |
| with_scale=False, |
| self.hidden_size, use_bias=False, name="linear_right" | ||
| ) | ||
| self.post_laurel_norm = Gemma3nRMSNorm( | ||
| # self.hidden_size, |
There was a problem hiding this comment.
This commented-out code should be removed to improve code cleanliness.
| "altup_active_idx": text_config["altup_active_idx"], | ||
| "altup_coef_clip": text_config.get("altup_coef_clip"), | ||
| "altup_correct_scale": text_config["altup_correct_scale"], | ||
| "rope_wavelength": rope_theta, # <-- ADD THIS |
There was a problem hiding this comment.
This comment seems to be a note for the developer and should be removed from the final version of the code.
| "rope_wavelength": rope_theta, # <-- ADD THIS | |
| "rope_wavelength": rope_theta, |
HF Reference
Closes #2320