Quickstart with Gemma on KerasNLP

Introduction

Gemma is a powerful language model that can be run using KerasNLP. This guide provides a quick demonstration of setting up and running Gemma, highlighting the necessary steps and configurations.

Access and Setup

To begin, ensure you have the necessary environment variables set up for accessing Kaggle datasets and models. This includes setting your KAGGLE_USERNAME and KAGGLE_KEY using the google.colab library.

import os
from google.colab import userdata
 
os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')
os.environ["KERAS_BACKEND"] = "jax"  # Or "tensorflow" or "torch".

Installation

Install the necessary dependencies to run Gemma on KerasNLP. Ensure you have the latest versions of keras-nlp and keras installed.

!pip install -U keras-nlp
!pip install -U keras==3.3.3

Quickstart

Once the setup is complete, you can proceed to download and initialize the Gemma model. The following code snippet demonstrates how to load the model and print its summary.

import keras_nlp
import keras
 
# Run at half precision.
keras.config.set_floatx("bfloat16")
 
# using 9B base model
gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma2_9b_en")
gemma_lm.summary()

Model Details

The Gemma model is extensive, with over 9 billion parameters, making it a robust tool for natural language processing tasks. Below is a summary of the model's architecture:

Conclusion

This quickstart guide provides a foundational understanding of how to set up and run the Gemma model using KerasNLP. With its extensive parameter set, Gemma is well-suited for a variety of NLP applications.

Reference: Google Gemini Gemma Cookbook by Google LLC.