Gemma 1.1 2B with Keras and KerasHub

Gemma 1.1 2B with Keras and KerasHub, Large language models have become an important part of modern artificial intelligence applications, but working with them does not always require a large and complicated software stack.

Gemma 1.1 2B, available through KerasHub, provides a relatively compact language model that developers can load, generate text with, and fine-tune using familiar Keras APIs.

This makes Gemma particularly interesting for developers and data scientists who already work with Python and Keras and want to experiment with generative AI.

The Keras implementation of Gemma 1.1 2B is an instruction-tuned model and is designed to work with Keras 3 across JAX, TensorFlow, and PyTorch backends. KerasHub currently lists gemma_1.1_instruct_2b_en as a 2-billion-parameter instruction-tuned model.

What Is Gemma 1.1 2B?

Gemma is a family of lightweight open models developed by Google. The Gemma 1.1 2B model is an updated version of the original instruction-tuned Gemma 2B model.

The “2B” designation refers to the approximately two-billion-parameter model class. KerasHub lists the actual model at around 2.51 billion parameters.

The important distinction is between the base Gemma model and the instruction-tuned version. KerasHub provides both gemma_2b_en and gemma_1.1_instruct_2b_en. The latter is the model used in this tutorial.

Because it is instruction tuned, it can be used for tasks where the user provides a prompt and expects a natural-language response.

Installing Keras and KerasHub

The first step is to install the required packages.

!pip install -U keras-hub
!pip install -U keras

Then import the required libraries:

import keras
import keras_hub
import numpy as np

KerasHub provides pretrained models through a convenient preset system, so developers do not need to manually download and assemble the model architecture and weights.

Loading Gemma 1.1 2B

The model can be loaded directly using GemmaCausalLM.from_preset().

gemma_lm = keras_hub.models.GemmaCausalLM.from_preset(
    "gemma_1.1_instruct_2b_en"
)

This loads the pretrained Gemma 1.1 instruction-tuned model.

KerasHub automatically provides the preprocessing required for normal string inputs when the model is created with from_preset().

That makes the basic workflow surprisingly simple: load the model, provide a prompt, and generate text.

Generating Text with Gemma

The generate() method is the main entry point for text generation.

gemma_lm.generate(
    "Keras is a",
    max_length=30
)

The max_length argument controls the length of the generated sequence.

You can also provide multiple prompts at once:

gemma_lm.generate(
    [
        "Keras is a",
        "Machine learning is"
    ],
    max_length=30
)

This is useful when experimenting with several prompts or processing multiple inputs in a batch. The current Keras documentation provides the same general generation workflow for Gemma models.

Controlling Text Generation

Language models can produce different results depending on the sampling strategy used during generation.

KerasHub allows the generation strategy to be changed using the compile() method.

For example:

gemma_lm.compile(sampler="top_k")

gemma_lm.generate(
    "Explain machine learning in simple terms.",
    max_length=100
)

KerasHub also supports beam search through BeamSampler.

gemma_lm.compile(
    sampler=keras_hub.samplers.BeamSampler(num_beams=2)
)

gemma_lm.generate(
    "Explain machine learning in simple terms.",
    max_length=100
)

The choice of sampler can affect the style and consistency of generated text. KerasHub uses greedy sampling by default and provides additional samplers for controlling generation.

Using Gemma for Data Science Tasks

A small language model such as Gemma can be useful for experimentation and application development.

For example, a data scientist could use it to generate explanations of statistical concepts:

prompt = """
Explain the difference between correlation and causation
using a simple example.
"""

response = gemma_lm.generate(
    prompt,
    max_length=150
)

print(response)

It can also be incorporated into applications that require natural-language generation, summarization, question answering, educational content, documentation assistance, and other text-based workflows.

However, generated responses should still be validated, particularly when the model is used for technical, financial, medical, or other high-impact information.

Fine-Tuning with Keras

One of the useful aspects of the Keras implementation is that Gemma can participate in the normal Keras training workflow.

For example:

features = [
    "Statistics helps us understand patterns in data.",
    "Machine learning can be used to build predictive models."
]

gemma_lm.fit(
    x=features,
    batch_size=2
)

This example demonstrates the training API, but it should not be interpreted as meaningful model fine-tuning with only two sentences. Real fine-tuning requires an appropriate dataset, training configuration, validation strategy, and sufficient computational resources.

KerasHub also provides support for parameter-efficient techniques such as LoRA, which can be used when adapting Gemma to specialized tasks without updating every model parameter.

Why Use Gemma with Keras?

The main attraction is simplicity.

Developers who already know Keras can use familiar concepts for loading models, generating predictions, batching inputs, and training. KerasHub also provides a consistent interface for accessing different pretrained models.

Another advantage is flexibility. The Keras implementation can run with different Keras 3 backends, including JAX, TensorFlow, and PyTorch.

This makes the combination useful for experimentation, learning, prototyping, and building custom generative AI applications.

Practical Considerations

Although Gemma 1.1 2B is smaller than many modern large language models, it still requires appropriate computational resources. Actual performance depends on the hardware, backend, precision, batch size, sequence length, and application.

It is also important to distinguish between a demonstration and a production system. Generating a response from a pretrained model is relatively straightforward, while building a reliable AI application requires additional work around evaluation, prompt design, data quality, monitoring, security, and output validation.

Conclusion

Gemma 1.1 2B provides a practical way for Python and Keras developers to explore generative AI without starting from a completely unfamiliar framework. With KerasHub, the model can be loaded from a preset, used for text generation, configured with different sampling strategies, and incorporated into Keras training workflows.

For data scientists, it offers an interesting bridge between traditional machine learning workflows and modern generative AI. Whether the goal is experimentation, education, prototyping, or developing a specialized application, Keras + KerasHub + Gemma 1.1 2B provides a straightforward starting point.