Safetensors conversion

[Bond099]

Description of the change

Reference

Colab Notebook

https://colab.research.google.com/drive/1naqf0sO2J40skndWbVMeQismjL7MuEjd?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.

[abheesht17]

Thanks for the PR, will take a look in a bit :)

[mattdangerw]

Thanks! Just left some initial comments.

[mattdangerw]

Let's add a unit test that calls this util and tries loading the result with transformers and seeing if it works. OK to add transformers to our ci environment here https://github.com/keras-team/keras-hub/blob/master/requirements-common.txt

[mattdangerw]

Nice! Please address the changes from the earlier PR as well

[abheesht17]

Thanks, nice work!

[abheesht17]

We should add the API export decorator here, similar to this: https://github.com/keras-team/keras-hub/blob/master/keras_hub/src/models/bloom/bloom_backbone.py#L15-L16

[abheesht17]

Also, do you think we should refactor some of the common code across models to a separate file? We can then expose that as the API.

So, this is how the directory keras_hub/src/utils/transformers/convert_to_safetensor/ will look like:

• export.py: this will have the common code. We will expose this as the API. This will also check if we support safetensor conversion for a given passed model yet.
• gemma.py: this will just have a way to create the weight dictionary for Gemma. Inside export.py, we will call the the weight conversion function specific to a specified model.

Pinging @mattdangerw to confirm if we should do this now or at a later point.

[mattdangerw]

I think we could land and do the API bit a later point. Though agree it's an important concern. I'm not sure if we want a method like model.save_to_preset() or a function like some_export(model). Any thoughts?

[Bond099]

I think structuring the export logic with a utility function (export_to_hf) and model-specific mappings (gemma.py) will enhance scalability and maintainability. New models can be added by creating a new file, while existing tests only need an import update.

[divyashreepathihalli]

+1 to Abheesht's comment we need an API instead of a script for Gemma, we already have that
https://github.com/keras-team/keras-hub/blob/master/tools/gemma/export_gemma_to_hf.py

[abheesht17]

Leaving comments since I don't see the changes we discussed last week.

[abheesht17]

Okay, reviewed. Let's fix the tests!

[abheesht17]

Why not just GemmaTokenizer?

[Bond099]

HF also imports the tokenizer as GemmaTokenizer… so just used KerasGemmaTokenizer to avoid confusion.

[abheesht17]

Oh, okay. Then it's better to call the HF import as HFGemmaTokenizer and the Keras one as GemmaTokenizer, I suppose

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

The code changes introduce a utility function to export Keras Gemma models to Hugging Face format, saving the configuration, weights, and tokenizer assets. The review focuses on improving the robustness of weight mapping, adding checks for empty weight dictionaries, and enhancing the warning message for missing vocabulary files.

[divyashreepathihalli]

lets make this a model-agnostic export utility. Rename file to safetensor_exporter.py
add a dict to maintain the mapping

MODEL_EXPORTERS = {
                   "GemmaBackbone": gemma_exporter.get_gemma_weights_map,
                   "LlamaBackbone": llama_exporter.get_llama_weights_map, # Future
}

and a user facing API function for the export

def export_to_safetensors(keras_model):
...

[divyashreepathihalli]

Implement exporter mapping for each model - for this PR's scope just the Gemma model that can serve as a prototype for other models

[abheesht17]

Yeah, let's land this PR first and do this in a separate PR: #2290 (comment)

[abheesht17]

LGTM, awesome work! Made a few cosmetic changes

[divyashreepathihalli]

Lets land the generic API in a different PR! Thanks for the great work!

[mattdangerw]

Thanks! Just a few changes.

[mattdangerw]

As a follow up (not on this PR), let's consider using keras_hub/src/tests/test_data/gemma_test_vocab.spm instead of retraining a new vocab here. Will be faster. Maybe leave a TODO?

[abheesht17]

I actually tried using it, and hit some issues. I didn't really dive deeper though, because I was lazy

[mattdangerw]

Actually let me switch to approval so I won't block merging this. But let's address these comment before merge!

[mattdangerw]

ty!