Gemma3 text keras hf checkpoint conversion #2433
Conversation
Summary of ChangesHello @kharshith-k, 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 checkpoint conversion tool by introducing the capability to export Keras models into the Hugging Face 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 adds a valuable feature to convert Keras Gemma3 models to the Hugging Face Safetensors format, enhancing interoperability. The implementation is comprehensive, covering configuration conversion, weight porting, and a validation step. I've provided a few suggestions to improve code clarity, maintainability, and adherence to the repository's style guide, primarily by improving docstrings, refactoring duplicated code, and ensuring deterministic validation.
| def convert_to_hf_config(keras_config): | ||
| """Convert Keras Gemma config to Hugging Face GemmaConfig.""" |
There was a problem hiding this comment.
The docstring for this function is missing Args and Returns sections, which is inconsistent with the repository's style guide. Providing detailed docstrings improves code clarity and maintainability.1
def convert_to_hf_config(keras_config):
"""Convert Keras Gemma config to Hugging Face GemmaConfig.
Args:
keras_config: A Keras Gemma3 config object from the backbone.
Returns:
A `transformers.Gemma3TextConfig` instance.
"""Style Guide References
Footnotes
-
The style guide requires all public functions to have Google-style docstrings, including comprehensive documentation for all parameters and return values. ↩
| def export_to_hf(backbone, keras_tokenizer, path): | ||
| """Convert a Keras Gemma model to Hugging Face format and save to path.""" |
There was a problem hiding this comment.
The docstring for this function is missing Args and Returns sections, which is inconsistent with the repository's style guide. Providing detailed docstrings improves code clarity and maintainability.1
def export_to_hf(backbone, keras_tokenizer, path):
"""Convert a Keras Gemma model to Hugging Face format and save to path.
Args:
backbone: A `keras_hub.models.Gemma3Backbone` instance.
keras_tokenizer: A `keras_hub.models.Gemma3Tokenizer` instance.
path: str. The path to save the Hugging Face model to.
"""Style Guide References
Footnotes
-
The style guide requires all public functions to have Google-style docstrings, including comprehensive documentation for all parameters and return values. ↩
| def to_torch(weight): | ||
| # Convert bfloat16 to float32 first, then to torch, then to bfloat16 | ||
| if hasattr(weight.dtype, "name") and "bfloat16" in str(weight.dtype): | ||
| weight = np.array(weight, dtype=np.float32) | ||
| return torch.from_numpy(weight).to(torch.bfloat16) |
There was a problem hiding this comment.
This helper function can be simplified to robustly handle various array types (like JAX arrays) and then used consistently throughout export_to_hf to reduce code duplication.
Currently, the conversion logic torch.from_numpy(np.array(weight, dtype=np.float32)).to(torch.bfloat16) is repeated for many weights. You can simplify to_torch to encapsulate this logic and improve maintainability.
With the suggested change, you can then refactor the rest of the function, for example:
q_kernel = block.attention.query_dense.get_weights()[0]
weights_dict[f"model.layers.{i}.self_attn.q_proj.weight"] = (
to_torch(q_kernel)
.permute(1, 0, 2)
.reshape(backbone.hidden_dim, -1)
.T
) def to_torch(weight):
# Convert array-like weights (e.g., from JAX) to a float32 NumPy
# array before creating a bfloat16 torch tensor for compatibility.
np_weight = np.array(weight, dtype=np.float32)
return torch.from_numpy(np_weight).to(torch.bfloat16)| top_k=50, | ||
| top_p=1.0, | ||
| ): | ||
| # Tokenize inpu |
There was a problem hiding this comment.
There is a typo in this comment.
# Tokenize input| outputs = model.generate( | ||
| **inputs, | ||
| max_new_tokens=max_new_tokens, | ||
| temperature=temperature, | ||
| top_k=top_k, | ||
| top_p=top_p, | ||
| do_sample=True, | ||
| ) |
There was a problem hiding this comment.
The infer function is used for a sanity check after converting the model. Currently, it uses sampling (do_sample=True), which makes the output non-deterministic. For a validation or sanity check step in a conversion script, it's better to use a deterministic generation strategy like greedy search to ensure consistent and predictable outputs.
outputs = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_k=top_k,
top_p=top_p,
do_sample=False,
)| "Export model to Safetensors format (HuggingFace-compatible). \ | ||
| Only for text-only models.", |
There was a problem hiding this comment.
The help string for the export_safetensors flag contains a backslash and extra whitespace, which will be literally included in the help message displayed to the user. It's better to use implicit string concatenation for multiline help messages to ensure clean formatting.
"Export model to Safetensors format (HuggingFace-compatible). "
"Only for text-only models.",
|
Thanks for the PR, the export to safetensors should be made available here https://github.com/keras-team/keras-hub/tree/master/keras_hub/src/utils/transformers/export.
|
4 participants
No description provided.