Changes to the Gemma3 backbone for Embedding Gemma model #2428
Conversation
Summary of ChangesHello @buildwithsuhana, 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 significantly enhances the Gemma3 backbone by introducing a dedicated embedding model. The core purpose is to provide a robust and accurate method for generating fixed-size sequence embeddings from the Gemma3 model's output, specifically by handling variable-length sequences through a new mean pooling mechanism that correctly accounts for padding. This change improves the model's utility for downstream tasks requiring dense vector representations of input sequences. Highlights
Using Gemini Code AssistThe 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 by creating a comment using either
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 Limitations & Feedback Gemini Code Assist 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a Gemma3EmbeddingModel and a MeanPooling layer to support embedding generation with Gemma 3 models. The changes are well-structured and include corresponding tests. My review focuses on ensuring adherence to the repository's style guide, particularly regarding documentation and serialization, to improve code quality and maintainability. I've suggested adding proper docstrings and get_config methods, and removing type hints from function signatures as per the contribution guidelines.
| return super().from_config(config) | ||
|
|
||
|
|
||
| class MeanPooling(layers.Layer): |
There was a problem hiding this comment.
Let's move this to gemma3_mean_pooling.py?
| super().__init__(**kwargs) | ||
| self.supports_masking = True | ||
|
|
||
| def call(self, inputs): |
There was a problem hiding this comment.
Instead of a tuple, can we just pass these arguments to call(...) separately instead of passing them as a tuple?
|
|
||
|
|
||
| @keras_hub_export("keras_hub.models.Gemma3Embedding") | ||
| class Gemma3EmbeddingModel(keras.Model): |
There was a problem hiding this comment.
I think we should just stick the layers on top in Gemma3Backbone instead of creating a separate class for this.
Can probably add an argument, named is_embedding_model, maybe. If is_embedding_model is True, add the dense layers + mean pooling layer, and return a dictionary, which returns the hidden states and the pooled output.
What do you think?
| ) | ||
| self.embedding_dim = embedding_dim | ||
|
|
||
| def call(self, inputs): |
There was a problem hiding this comment.
We don't generally have call(...) for KerasHub models because these models are Functional models.
| self.projection_layer = layers.Dense( | ||
| embedding_dim, dtype=backbone.dtype, name="embedding_projection" |
There was a problem hiding this comment.
Is there just one dense layer for Embedding Gemma? I thought there were several.
This PR introduces the new Gemma 3 backbone for embedding gemma model. It includes:
MeanPooling Layer: This new layer correctly computes the average of token embeddings while ignoring padded tokens using a mask.
The Gemma3EmbeddingModel: This model integrates a Gemma3Backbone, your new MeanPooling layer, and a final Dense projection layer to produce fixed-size sequence embeddings. It is also fully serializable.