Hugging Face and KerasHub Integration

Hugging Face and KerasHub have introduced a shared model save format that allows KerasHub users to directly load models created with the Transformers library from the Hugging Face Hub. This integration removes the previous limitation where KerasHub users could only access models specifically built for KerasHub, granting them access to a library of over 300,000 pre-trained models.

Direct Access to Transformers Models

KerasHub can now load Transformers library checkpoints directly using the from_preset method. This enables users to utilize a vast array of fine-tuned models that were not originally created with Keras.

Initially, this integration supports the following architectures:

  • Gemma (versions 1 and 2)
  • Llama 3
  • PaliGemma

Multi-Framework Deployment

Because KerasHub models can operate using TensorFlow, JAX, or PyTorch backends, this integration allows users to load a Hugging Face checkpoint into any of these frameworks with a single line of code. This capability simplifies the process of porting models for specific needs, such as deploying to TFLite for serving or utilizing JAX for research purposes.

Technical Implementation

The integration works by mapping configuration variables, weight names, and tokenizer vocabularies between the two libraries. Since Transformers models are stored as JSON configuration files, tokenizer files, and safetensors weights, KerasHub can create a compatible checkpoint as long as both libraries have the modeling code for the relevant architecture. This conversion process is handled internally by the libraries, requiring no manual conversion from the user.

Usage and Configuration

To use the integration, users must update to keras-hub and keras>=3.3.3.

Text Generation

Users can load a Transformers model and generate text using the .generate method. For example, loading a Llama 3 model from the Hub:

from keras_hub.models import Llama3CausalLM

causal_lm = Llama3CausalLM.from_preset(
    "hf://NousResearch/Hermes-2-Pro-Llama-3-8B"
)

prompts = ["Your prompt here"]
causal_lm.generate(prompts, max_length=200)

Precision and Backend Control

KerasHub allows for easy adjustment of model precision and the underlying compute backend:

  • Changing Precision: Precision can be set via keras.config.set_dtype_policy("bfloat16") before loading the model.
  • Switching Backends: By setting the environment variable os.environ["KERAS_BACKEND"] = "jax", users can run the loaded Transformers checkpoint using the JAX backend.

Supported Models

Beyond Llama 3, the integration explicitly supports:

  • Gemma 2: Users can load Gemma 2 models (e.g., google/gemma-2-9b) directly.
  • PaliGemma: Any PaliGemma safetensor checkpoint, including fine-tuned versions, can be integrated into a KerasHub pipeline.

Sources