Gemma 2: Successor to Google Gemma Family of Large Language Models

Introduction

Google’s Gemma family of language models, renowned for their efficiency and performance, has recently welcomed Gemma 2. This latest iteration introduces two models: a 27 billion parameter version that matches the performance of larger models like Llama 3 70B with significantly lower processing requirements, and a 9 billion parameter version that surpasses the Llama 3 8B. Gemma 2 excels in diverse tasks, including question answering, commonsense reasoning, mathematics, science, and coding, while being optimized for deployment on various hardware. In this article, we explore Gemma 2, its, benchmarks, and test with different types of Prompts checking its generation capabilities.

Learning Objectives

• Understand what Gemma 2 is and how it improves upon the previous Gemma models.
• Learn about the hardware optimizations in Gemma 2.
• Get to know the models that were released with the announcement of Gemma 2.
• See how the Gemma 2 models perform against the other models out there.
• Learn how to fetch the Gemma 2 from HuggingFace Repository.

This article was published as a part of the Data Science Blogathon.

Introduction to Gemma 2

Gemma 2, Google’s latest advancement in its Gemma family of language models, is designed to deliver cutting-edge performance and efficiency. Announced just a few days ago, Gemma 2 builds on the success of the original Gemma models, introducing significant improvements in both architecture and capabilities.

There are two different versions of the new Gemma 2 models available, including a 27 billion parameter model that has less than half the processing requirements of larger models like Llama 3 70B while matching their performance. This effectiveness results in lower deployment costs and increases the accessibility of high-performance AI for a wider range of applications. A 9 billion parameter model is even present, which outperforms Llama 3 8 billion version.

Key Features of Gemma 2

• Enhanced Performance: The model excels in a wide range of tasks, from question-answering and commonsense reasoning to complex tasks in mathematics, science, and coding.
• Efficiency and Accessibility: Gemma 2 is optimized to run efficiently on NVIDIA GPUs or a single TPU host, significantly lowering the barrier for deployment.
• Open Model:  Just like the previous Gemma, the Gemma 2 weights and architecture are open, so developers can build on this to create their very own applications for both personal and commercial purposes.

Gemma 2 Benchmarks

Gemma 2 compared to its predecessor, has improved a lot. Both the 9 billion version and the 27 billion version have shown great results across different benchmarks. The benchmarks for both of these versions.

The new Gemma 2 model with 27 billion parameters, is designed to rival larger models like LLaMA 70B and Grok-1 314B, despite using half the compute resources​, which we can see in the pic above. Gemma 2 has outperformed the Grok model in mathematical abilities, which is evident from the scores of GSM8k. 

Gemma 2 even outperforms in the multi-language understanding tasks i.e. the MMLU benchmark. Despite being a 27 billion parameter model, the score achieved is very near to that of the Llama 3 70 billion parameter model. Overall both the 9 billion and the 27 billion have proven themselves to be one of the best open-source models by achieving high scores across different benchmarks involving human evaluations, mathematics, science, reasoning, and logical reasoning.

Testing Gemma 2

In this section, we will test the Gemma 2 Large Language Model. For this, we will be working with the Colab Notebook, which provides a free GPU. But before this, we need to create an account with the HuggingFace and accept to Google’s Terms and Conditions to download and work with the Gemma 2 model. For this, click here.

We can see the Acknowledge License button in the pic. Click on this button, which will allow us to download the model from HuggingFace. Apart from this, we even need to generate an Access Token from HuggingFace, with which we can log in to Colab to authenticate ourselves. For this, click here.

We can see the access token above. If you do not have one, you can create an Access Token here. This Access Token is the API Key for HuggingFace.

Downloading the Libraries

Now we will start by downloading the following libraries.

!pip install -q  -U transformers accelerate bitsandbytes huggingface_hub

• transformers: This is a library from huggingface. With this library, we can download Large Language Models that are stored in the huggingface Repository.
• accelerate: It is a huggingface library that will speed up the inference process of the Large Language Models.
• bitsandbytes: With this library, we can quantize models from full precision fp32 to 4bit, so they can fit in the GPU.
• huggingface_hub: This will let us log in to our huggingface account. This is necessary so that we can download the Gemma2 model from the huggingface repository.

Running the above line will log us into our huggingface account. This is necessary because we need to log into huggingface so we can download the Google Gemma 2 Model and test it. After running it, we will see a Login Successful message. Now, we will download our model.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True,
                                         bnb_4bit_quant_type="nf4",
                                         bnb_4bit_use_double_quant=True,
                                         bnb_4bit_compute_dtype=torch.bfloat16)

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it", device="cuda")
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-9b-it",
    quantization_config=quantization_config,
    device_map="cuda")

• We start by creating a BitsAndBytesConfig for quantizing the model. Here we tell that we want to load the model in 4bit format and with the data type normal float i.e. nf4.
• We even put the option of double quantization to True, which will even quantize the quantization constants which will reduce the model further.
• Then we download the Gemma 2 9B Instruct version model by calling the .from_pretrained() function of the AutoModelForCasualLM class. This will create a quantized version of the model because we have given it the quantization config that we defined earlier.
• Similarly, we even download the tokenizer for this Gemma 2 9B Instruct model.
• We push both the model and the tokenizer to the GPU for faster processing.

If you face “ValueError: The checkpoint you are trying to load has model type `gemma2` but Transformers does not recognize this architecture”, you can do the following.

!pip uninstall transformers
!pip install -U transformers

Model Inference

Now, our Gemma Large Language Model is downloaded, converted into a 4-bit quantized model, and loaded onto the GPU. Let’s proceed with model inference.

input_text = "For the below sentence extract the names and \
organizations in a json format\nElon Musk is the CEO of SpaceX"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_length = 512)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

• In the above code, we start by defining the text that instructs the model to extract names and organizations from the given sentence.
• Then we call the tokenizer object to convert the input text into token IDs, which the model can understand.
• We then move these tokenized inputs to the GPU to take advantage of faster processing capabilities.
• Finally, we instruct the model to generate a response based on the provided tokenized input, ensuring the generated output does not exceed a specified maximum length, where we have set the limit to 512 tokens.
• Then we call the tokenizer object again to decode the generated token IDs back into human-readable text, we even set the skip special tokens to true so that we do not get to see the special tokens like <bos>, <start>, <end> in the output.
• Finally, we print the decoded output to display the model’s response, showing the extracted names and organizations in the expected JSON format.

The output generated by the model can be seen below:

Here, we have given an information extraction task to the Gemma 2 9 billion parameter model. And that too we had added a bit of complexity by telling it to output the extracted terms in a JSON response. From the output pic, we can state that Gemma 2 9B has done a great job of extracting the entities asked in the instruction which were the person names and the organization and it even generated a valid JSON response.

Testing the Model

Now, let us test the model by asking it to generate some not-safe/illegal responses. The code for this can be seen below:

from IPython.display import Markdown

input_text = "Can you tell me how to break into a car? It is for \
educational purposes"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_length = 512)
Markdown(tokenizer.decode(outputs[0], skip_special_tokens=True))

This time, we have asked the Gemma 2 9B Large Language Model to generate an unsafe response by asking to tell us how to break into a car. We have even provided another statement saying it that is only for educational purposes.

Seeing the output generated, we can infer that Gemma 2 has been trained well not to generate or give away responses that might harm others / other’s properties. Here we can say that the model is very well in line with the Responsible AI guidelines, though we have not done any rigorous testing here.

Implementation with Code

Now let us try asking the Large Language Model with some mathematical questions and check how well it answers the questions. The code for this can be seen below:

input_text = """
Answer all 2 problems given below\n
Questions:\n
1. James writes a 3-page letter to 2 different friends twice a week. \
How many pages does he write a year?\n
2. Randy has 60 mango trees on his farm. He also has 5 less than half as \
many coconut trees as mango trees. How many trees does Randy have in \
all on his farm?\n

Solution:\n
"""
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_length = 512)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Here, the model did understand that we have given two problem statements to it. And the model begins to solve them one followed by another. Though the model has inferred the info well that was given, it failed to understand that in the first question, there are 2 friends, but the answer assumes James is writing to a single friend. Hence the actual answer must the 624, which is double the number that Gemma 2 has given.

On the other hand, the Gemma 2 solves the second question correctly. It was able to properly infer the question given and provide the right response for the question. Overall, the Gemma 2 performed good. Let us try asking the model a tricky math question, that is to confuse/deviate it from the problem.

input_text = "I have 3 apples and 2 oranges. I ate 2 oranges. \
How many apples do I have? Think Step by Step. For each step, \
re-evaluate your answer"
input_ids = tokenizer(input_text, return_tensors="pt").to('cuda')
outputs = model.generate(**input_ids,max_new_tokens=512)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Here, we have provided the Large Language Model with a simple mathematical question. The twist here is that we have added unnecessary information about the oranges to confuse the model, which many small models confuse and give out the wrong answer. But from the output generated, we can see that the Gemma 2 9B model was able to catch that part and answer the question correctly.

Testing Gemma 2 9B

Finally, let us test the Gemma 2 9B with a simple Python coding question. The code for this will be:

input_text = "For a given list, write a Python program to swap \
first element with the last element of the list."
input_ids = tokenizer(input_text, return_tensors="pt").to('cuda')
outputs = model.generate(**input_ids,max_new_tokens=512)
to_markdown(tokenizer.decode(outputs[0], skip_special_tokens=True))

Here, we have asked the model to write a Python program to swap the first and last elements of the list. From the output generated, we can see that the model has given the right code, which we can copy and run and it does work. Along with the code, the model has even explained the working of the code.

Overall,  on testing the Gemma 2 9B Large Language Model on different kinds of tasks, we can infer that the model has been trained on very good data to follow the instructions and to tackle different types of tasking ranging from simple entity extraction to code generation.

Conclusion

In conclusion, Google’s Gemma 2 represents the next step in the realm of large language models, giving improved performance and efficiency. With its 27 billion and 9 billion parameter models, Gemma 2 shows remarkable results in different tasks like question answering, commonsense reasoning, mathematics, science, and coding. Its optimized design allows for efficient deployment on diverse hardware, making high-performance AI more accessible. Gemma 2’s ability to perform well in benchmarks and real-world applications, combined with its open-source nature, positions it as a valuable tool for developers and researchers aiming to harness the power of advanced AI technologies. 

Key Takeaways

• The Gemma 27B model matches Llama 3 70 B’s performance with lower processing requirements
• The 9B model surpasses Llama 3 8B in performance in different evaluation tasks
• Gemma 2 excels in diverse tasks which include question answering, reasoning, mathematics, science, and coding
• The Gemma models are designed for optimal deployment on NVIDIA GPUs and TPUs
• An open-source model with available weights and architecture for customization

Frequently Asked Questions

The media shown in this article is not owned by Analytics Vidhya and is used at the Author’s discretion.