Add generic export_to_transformers to the base classes #2346

Bond099 · 2025-07-26T04:35:29Z

Description of the change

Reference

Colab Notebook

https://colab.research.google.com/drive/1CNUkqbRTBPirTaU1-2UrWuK2dTZhgOtn?usp=sharing

Checklist

I have added all the necessary unit tests for my change.
I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
My PR is based on the latest changes of the main branch (if unsure, rebase the code).
I have followed the Keras Hub Model contribution guidelines in making these changes.
I have followed the Keras Hub API design guidelines in making these changes.
I have signed the Contributor License Agreement.

gemini-code-assist

Summary of Changes

Hello @Bond099, I'm Gemini Code Assist¹! 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 introduces a significant enhancement by providing direct export capabilities from Keras Hub models and tokenizers to the Hugging Face Transformers format. This change streamlines the process of converting Keras-based components for use within the Hugging Face ecosystem, improving interoperability and ease of use for developers working with both frameworks.

Highlights

Direct Hugging Face Export for Models: I've added a new export_to_transformers method to the Backbone and CausalLM base classes. This allows users to directly export their Keras models (or just the backbone) into a format compatible with Hugging Face Transformers, including configuration and weights, simplifying interoperability.
Direct Hugging Face Export for Tokenizers: A corresponding export_to_transformers method has been added to the Tokenizer base class. This enables the direct export of Keras tokenizers, saving their assets in a format that Hugging Face Transformers can readily consume.
Refactored Export Logic: The core export logic in keras_hub/src/utils/transformers/export/hf_exporter.py has been modularized. The existing export_to_safetensors function now accepts a verbose argument, and two new functions, export_backbone and export_tokenizer, have been introduced to handle the specific export needs of model backbones and tokenizers, respectively.
Gemma Model Export Improvements: For Gemma models, I've updated the Hugging Face configuration to explicitly set tie_word_embeddings to True. Additionally, the explicit weight tying for lm_head.weight in the Gemma exporter has been removed, as this is now implicitly handled by the tie_word_embeddings configuration in Hugging Face.
Updated Export Tests: The gemma_test.py file has been updated to reflect and validate the new export_to_transformers methods. Tests now separately export and verify the backbone, tokenizer, and the full CausalLM model, ensuring that all components are correctly converted and loadable by Hugging Face.

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 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 or fill out our survey 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.

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. ↩

gemini-code-assist

Code Review

This pull request introduces a generic export_to_transformers method to the Backbone, CausalLM, and Tokenizer base classes, enabling models to be exported to the Hugging Face Transformers format. The implementation adds new exporter functions and updates tests accordingly. I've provided feedback to reduce code duplication, improve error message clarity, and enhance code documentation.

keras_hub/src/utils/transformers/export/hf_exporter.py

keras_hub/src/models/causal_lm.py

divyashreepathihalli · 2025-07-29T23:42:18Z

Can you add a colab example demo?

keras_hub/src/utils/transformers/export/hf_exporter.py

mattdangerw

Thanks! Overall looks good! Just had some questions on the API and a request for more testing.

keras_hub/src/models/backbone.py

keras_hub/src/models/causal_lm.py

keras_hub/src/models/backbone.py

abheesht17

Thanks! LGTM, mostly. Left one important question and a few nits.

keras_hub/src/models/backbone.py

keras_hub/src/models/causal_lm.py

keras_hub/src/models/causal_lm_preprocessor.py

keras_hub/src/models/backbone.py

keras_hub/src/utils/transformers/export/gemma.py

keras_hub/src/utils/transformers/export/hf_exporter.py

…neric_export

abheesht17

LGTM, thanks! I'm going to let @mattdangerw do another pass on this

mattdangerw

Thanks!

keras_hub/src/models/backbone_test.py

keras_hub/src/models/causal_lm.py

keras_hub/src/models/causal_lm_preprocessor.py

keras_hub/src/models/causal_lm_preprocessor_test.py

keras_hub/src/utils/transformers/export/hf_exporter.py

…port

…neric_export

mattdangerw

One last nit!

mattdangerw · 2025-08-15T20:54:59Z

keras_hub/src/models/task_test.py

+
+    def _create_gemma_for_export_tests(self):
+        proto = os.path.join(
+            os.path.dirname(__file__),


Can't we just do this the same as our other usages?

proto = os.path.join(self.get_test_data_dir(), "gemma_export_vocab.spm")

mattdangerw · 2025-08-15T20:55:16Z

keras_hub/src/models/causal_lm_preprocessor_test.py

@@ -43,3 +49,26 @@ def test_from_preset_errors(self):
        with self.assertRaises(ValueError):
            # No loading on a non-keras model.
            GPT2CausalLMPreprocessor.from_preset("hf://spacy/en_core_web_sm")
+
+    def test_export_supported_preprocessor(self):
+        proto = os.path.join(


Can't we just do this the same as our other usages?

proto = os.path.join(self.get_test_data_dir(), "gemma_export_vocab.spm")

mattdangerw · 2025-08-15T20:55:26Z

keras_hub/src/tokenizers/tokenizer_test.py

@@ -113,3 +114,32 @@ def test_save_to_preset(self, cls, preset_name, tokenizer_type):
        # Check config class.
        tokenizer_config = load_json(save_dir, TOKENIZER_CONFIG_FILE)
        self.assertEqual(cls, check_config_class(tokenizer_config))
+
+    def test_export_supported_tokenizer(self):
+        proto = os.path.join(


proto = os.path.join(self.get_test_data_dir(), "gemma_export_vocab.spm")

mattdangerw · 2025-08-15T20:55:51Z

keras_hub/src/utils/transformers/export/gemma_test.py

-            user_defined_symbols=["<start_of_turn>", "<end_of_turn>"],
-        )
-        tokenizer = GemmaTokenizer(proto=f"{proto_prefix}.model")
+        proto = os.path.join(


proto = os.path.join(self.get_test_data_dir(), "gemma_export_vocab.spm")

gemini-code-assist bot reviewed Jul 26, 2025

View reviewed changes

keras_hub/src/utils/transformers/export/hf_exporter.py Outdated Show resolved Hide resolved

keras_hub/src/utils/transformers/export/hf_exporter.py Outdated Show resolved Hide resolved

keras_hub/src/models/causal_lm.py Outdated Show resolved Hide resolved

mattdangerw added the kokoro:force-run Runs Tests on GPU label Jul 31, 2025

kokoro-team removed the kokoro:force-run Runs Tests on GPU label Jul 31, 2025

abheesht17 reviewed Jul 31, 2025

View reviewed changes

keras_hub/src/utils/transformers/export/hf_exporter.py Outdated Show resolved Hide resolved

mattdangerw reviewed Aug 5, 2025

View reviewed changes

keras_hub/src/models/backbone.py Outdated Show resolved Hide resolved

keras_hub/src/models/causal_lm.py Outdated Show resolved Hide resolved

keras_hub/src/models/backbone.py Outdated Show resolved Hide resolved

abheesht17 added the kokoro:force-run Runs Tests on GPU label Aug 7, 2025

kokoro-team removed the kokoro:force-run Runs Tests on GPU label Aug 7, 2025

abheesht17 reviewed Aug 7, 2025

View reviewed changes

Bond099 added 5 commits August 8, 2025 18:52

export_to_transformers

37f33e1

Address comments

9a17cde

Fixes AutoTokenizer and AutoModel compatibility

f12c998

Add tests

cd4bdff

Add lm_head

ba30d5e

Bond099 force-pushed the generic_export branch from d8ae4a9 to ba30d5e Compare August 8, 2025 13:35

Bond099 added 7 commits August 8, 2025 20:19

export_to_transformers

21ebaa3

Address comments

6fa66b0

Fixes AutoTokenizer and AutoModel compatibility

b9c23d6

Add tests

fefa0b8

Add lm_head

98d61fa

Merge branch 'generic_export' of github.com:Bond099/keras-hub into ge…

399c456

…neric_export

resolve gemma_test

132178f

abheesht17 approved these changes Aug 12, 2025

View reviewed changes

mattdangerw reviewed Aug 14, 2025

View reviewed changes

Bond099 added 5 commits August 14, 2025 18:55

Merge branch 'keras-team:master' into generic_export

3a3d0b7

Merge branch 'master' of github.com:Bond099/keras-hub into generic_ex…

8fa07c3

…port

Address comments and add llama model export support

031c847

Merge branch 'generic_export' of github.com:Bond099/keras-hub into ge…

9356d34

…neric_export

Move llama to a new PR

1d51b97

mattdangerw approved these changes Aug 15, 2025

View reviewed changes

mattdangerw added the kokoro:force-run Runs Tests on GPU label Aug 15, 2025

kokoro-team removed the kokoro:force-run Runs Tests on GPU label Aug 15, 2025

Add generic export_to_transformers to the base classes #2346

Are you sure you want to change the base?

Add generic export_to_transformers to the base classes #2346

Uh oh!

Conversation

Bond099 commented Jul 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description of the change

Reference

Colab Notebook

Checklist

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

divyashreepathihalli commented Jul 29, 2025

Uh oh!

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

abheesht17 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

abheesht17 left a comment

Choose a reason for hiding this comment

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

mattdangerw Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

mattdangerw Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

mattdangerw Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

mattdangerw Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Bond099 commented Jul 26, 2025 •

edited

Loading