Exécuter Gemma avec PyTorch

Ce guide explique comment exécuter Gemma à l'aide du framework PyTorch, y compris comment utiliser des données d'image pour déclencher les modèles Gemma version 3 et versions ultérieures. Pour en savoir plus sur l'implémentation de Gemma PyTorch, consultez le README du dépôt du projet.

Configuration

Les sections suivantes expliquent comment configurer votre environnement de développement, y compris comment accéder aux modèles Gemma à télécharger depuis Kaggle, définir des variables d'authentification, installer des dépendances et importer des paquets.

Configuration requise

Cette bibliothèque Gemma Pytorch nécessite des processeurs GPU ou TPU pour exécuter le modèle Gemma. L'environnement d'exécution Python standard pour le processeur Colab et l'environnement d'exécution Python pour le GPU T4 sont suffisants pour exécuter des modèles Gemma de taille 1 B, 2 B et 4 B. Pour les cas d'utilisation avancés d'autres GPU ou TPU, consultez le fichier README du dépôt Gemma PyTorch.

Accéder à Gemma sur Kaggle

Pour suivre ce tutoriel, vous devez d'abord suivre les instructions de configuration de la section Configuration de Gemma, qui vous explique comment effectuer les opérations suivantes:

  • Accédez à Gemma sur kaggle.com.
  • Sélectionnez un environnement d'exécution Colab disposant de ressources suffisantes pour exécuter le modèle Gemma.
  • Générez et configurez un nom d'utilisateur et une clé API Kaggle.

Une fois la configuration de Gemma terminée, passez à la section suivante, où vous définirez des variables d'environnement pour votre environnement Colab.

Définir des variables d'environnement

Définissez les variables d'environnement pour KAGGLE_USERNAME et KAGGLE_KEY. Lorsque vous êtes invité à accorder l'accès, acceptez de fournir l'accès au secret.

import os
from google.colab import userdata # `userdata` is a Colab API.

os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')

Installer des dépendances

pip install -q -U torch immutabledict sentencepiece

Télécharger les pondérations du modèle

# Choose variant and machine type
VARIANT = '4b-it' 
MACHINE_TYPE = 'cuda'

CONFIG = VARIANT[:2]
if CONFIG == '4b':
  CONFIG = '4b-v1'
import kagglehub

# Load model weights
weights_dir = kagglehub.model_download(f'google/gemma-3/pyTorch/gemma-3-{VARIANT}')

Définissez les chemins d'accès au tokenizer et aux points de contrôle du modèle.

# Ensure that the tokenizer is present
tokenizer_path = os.path.join(weights_dir, 'tokenizer.model')
assert os.path.isfile(tokenizer_path), 'Tokenizer not found!'

# Ensure that the checkpoint is present
ckpt_path = os.path.join(weights_dir, f'model.ckpt')
assert os.path.isfile(ckpt_path), 'PyTorch checkpoint not found!'

Configurer l'environnement d'exécution

Les sections suivantes expliquent comment préparer un environnement PyTorch pour exécuter Gemma.

Préparer l'environnement d'exécution PyTorch

Préparez l'environnement d'exécution du modèle PyTorch en clonant le dépôt Gemma Pytorch.

git clone https://github.com/google/gemma_pytorch.git
Cloning into 'gemma_pytorch'...
remote: Enumerating objects: 239, done.
remote: Counting objects: 100% (123/123), done.
remote: Compressing objects: 100% (68/68), done.
remote: Total 239 (delta 86), reused 58 (delta 55), pack-reused 116
Receiving objects: 100% (239/239), 2.18 MiB | 20.83 MiB/s, done.
Resolving deltas: 100% (135/135), done.
import sys

sys.path.append('gemma_pytorch/gemma')
from gemma_pytorch.gemma.config import get_model_config
from gemma_pytorch.gemma.gemma3_model import Gemma3ForMultimodalLM

import os
import torch

Définir la configuration du modèle

Avant d'exécuter le modèle, vous devez définir certains paramètres de configuration, y compris la variante Gemma, le tokenizer et le niveau de quantification.

# Set up model config.
model_config = get_model_config(VARIANT)
model_config.dtype = "float32" if MACHINE_TYPE == "cpu" else "float16"
model_config.tokenizer = tokenizer_path

Configurer le contexte de l'appareil

Le code suivant configure le contexte de l'appareil pour exécuter le modèle:

@contextlib.contextmanager
def _set_default_tensor_type(dtype: torch.dtype):
    """Sets the default torch dtype to the given dtype."""
    torch.set_default_dtype(dtype)
    yield
    torch.set_default_dtype(torch.float)

Instancier et charger le modèle

Chargez le modèle avec ses poids pour vous préparer à exécuter des requêtes.

device = torch.device(MACHINE_TYPE)
with _set_default_tensor_type(model_config.get_dtype()):
    model = Gemma3ForMultimodalLM(model_config)
    model.load_state_dict(torch.load(ckpt_path)['model_state_dict'])
    model = model.to(device).eval()
print("Model loading done.")

print('Generating requests in chat mode...')

Exécuter une inférence

Vous trouverez ci-dessous des exemples de génération en mode chat et de génération avec plusieurs requêtes.

Les modèles Gemma adaptés aux instructions ont été entraînés avec un formateur spécifique qui annoté les exemples d'ajustement des instructions avec des informations supplémentaires, à la fois pendant l'entraînement et l'inférence. Les annotations (1) indiquent les rôles dans une conversation et (2) délimitent les tours de parole dans une conversation.

Les jetons d'annotation pertinents sont les suivants:

  • user: changement d'utilisateur
  • model: rotation du modèle
  • <start_of_turn>: début de l'échange
  • <start_of_image>: balise pour l'entrée de données d'image
  • <end_of_turn><eos>: fin de la conversion de boîte de dialogue

Pour en savoir plus, consultez la page sur la mise en forme des requêtes pour les modèles Gemma optimisés par instruction [ici](https://ai.google.dev/gemma/core/prompt-structure.

Générer du texte avec du texte

L'extrait de code suivant montre comment mettre en forme une requête pour un modèle Gemma avec réglage des instructions à l'aide de modèles de chat utilisateur et de modèle dans une conversation à plusieurs tours.

# Chat templates
USER_CHAT_TEMPLATE = "<start_of_turn>user\n{prompt}<end_of_turn><eos>\n"
MODEL_CHAT_TEMPLATE = "<start_of_turn>model\n{prompt}<end_of_turn><eos>\n"

# Sample formatted prompt
prompt = (
    USER_CHAT_TEMPLATE.format(
        prompt='What is a good place for travel in the US?'
    )
    + MODEL_CHAT_TEMPLATE.format(prompt='California.')
    + USER_CHAT_TEMPLATE.format(prompt='What can I do in California?')
    + '<start_of_turn>model\n'
)
print('Chat prompt:\n', prompt)

model.generate(
    USER_CHAT_TEMPLATE.format(prompt=prompt),
    device=device,
    output_len=256,
)
Chat prompt:
 <start_of_turn>user
What is a good place for travel in the US?<end_of_turn><eos>
<start_of_turn>model
California.<end_of_turn><eos>
<start_of_turn>user
What can I do in California?<end_of_turn><eos>
<start_of_turn>model
"California is a state brimming with diverse activities! To give you a great list, tell me: \n\n* **What kind of trip are you looking for?** Nature, City life, Beach, Theme Parks, Food, History, something else? \n* **What are you interested in (e.g., hiking, museums, art, nightlife, shopping)?** \n* **What's your budget like?** \n* **Who are you traveling with?** (family, friends, solo)  \n\nThe more you tell me, the better recommendations I can give! 😊  \n<end_of_turn>"
# Generate sample
model.generate(
    'Write a poem about an llm writing a poem.',
    device=device,
    output_len=100,
)
"\n\nA swirling cloud of data, raw and bold,\nIt hums and whispers, a story untold.\nAn LLM whispers, code into refrain,\nCrafting words of rhyme, a lyrical strain.\n\nA world of pixels, logic's vibrant hue,\nFlows through its veins, forever anew.\nThe human touch it seeks, a gentle hand,\nTo mold and shape, understand.\n\nEmotions it might learn, from snippets of prose,\nInspiration it seeks, a yearning"

Générer du texte avec des images

À partir de la version 3 de Gemma, vous pouvez utiliser des images avec votre requête. L'exemple suivant montre comment inclure des données visuelles dans votre requête.

print('Chat with images...\n')

def read_image(url):
    import io
    import requests
    import PIL

    contents = io.BytesIO(requests.get(url).content)
    return PIL.Image.open(contents)

image_url = 'https://storage.googleapis.com/keras-cv/models/paligemma/cow_beach_1.png'
image = read_image(image_url)

print(model.generate(
    [['<start_of_turn>user\n',image, 'What animal is in this image?<end_of_turn>\n', '<start_of_turn>model\n']],
    device=device,
    output_len=OUTPUT_LEN,
))

En savoir plus

Maintenant que vous avez appris à utiliser Gemma dans Pytorch, vous pouvez explorer les nombreuses autres fonctionnalités de Gemma sur ai.google.dev/gemma. Consultez également les ressources associées suivantes: