Ver en ai.google.dev | Ejecutar en Google Colab | Abrir en Vertex AI | Ver código fuente en GitHub |
Descripción general
Gemma es una familia de modelos grandes de lenguaje, ligeros y de vanguardia, basados en la investigación y tecnología de Google DeepMind Gemini. En este instructivo, se muestra cómo realizar inferencias o muestras básicas con el modelo Gemma 2B Instruct a través de la biblioteca gemma
de Google DeepMind que se escribió con JAX (una biblioteca de computación numérica de alto rendimiento), Flax (la biblioteca de redes neuronales basada en JAX), Orbax (una biblioteca basada en JAXizer/detoken1 para las utilidades de entrenamiento como token de control) y SentencePiece Aunque Flax no se usa directamente en este cuaderno, se usó para crear Gemma.
Este notebook se puede ejecutar en Google Colab con la GPU T4 gratuita (ve a Editar > Configuración del notebook > En Acelerador de hardware, selecciona GPU T4).
Instalar
1. Configurar el acceso a Kaggle para Gemma
Para completar este instructivo, primero debes seguir las instrucciones de configuración en la configuración de Gemma, que te muestran cómo hacer lo siguiente:
- Obtén acceso a Gemma en kaggle.com.
- Selecciona un entorno de ejecución de Colab con recursos suficientes para ejecutar el modelo de Gemma.
- Genera y configura un nombre de usuario y una clave de API Kaggle.
Después de completar la configuración de Gemma, continúa con la siguiente sección, en la que establecerás variables de entorno para tu entorno de Colab.
2. Configura las variables de entorno
Configura las variables de entorno para KAGGLE_USERNAME
y KAGGLE_KEY
. Cuando aparezca el mensaje "¿Quieres otorgar acceso?", acepta proporcionar acceso al secreto.
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')
3. Cómo instalar la biblioteca gemma
Este notebook se enfoca en el uso de una GPU de Colab gratuita. Para habilitar la aceleración de hardware, haz clic en Editar > Configuración del notebook > selecciona GPU T4 > Guardar.
A continuación, debes instalar la biblioteca gemma
de Google DeepMind desde github.com/google-deepmind/gemma
. Si recibes un error sobre el "agente de resolución de dependencia de pip", generalmente puedes ignorarlo.
pip install -q git+https://github.com/google-deepmind/gemma.git
Carga y prepara el modelo de Gemma
- Carga el modelo de Gemma con
kagglehub.model_download
, que toma tres argumentos:
handle
: El controlador del modelo de Kagglepath
: Es la ruta local (cadena opcional).force_download
: Fuerza a volver a descargar el modelo (booleano opcional).
GEMMA_VARIANT = '2b-it' # @param ['2b', '2b-it'] {type:"string"}
import kagglehub
GEMMA_PATH = kagglehub.model_download(f'google/gemma/flax/{GEMMA_VARIANT}')
Downloading from https://www.kaggle.com/api/v1/models/google/gemma/flax/2b-it/2/download... 100%|██████████| 3.67G/3.67G [00:35<00:00, 110MB/s] Extracting model files...
print('GEMMA_PATH:', GEMMA_PATH)
GEMMA_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2
- Verifica la ubicación de las ponderaciones del modelo y del tokenizador, y luego configura las variables de la ruta de acceso. El directorio del tokenizador estará en el directorio principal en el que descargaste el modelo, mientras que los pesos del modelo estarán en un subdirectorio. Por ejemplo:
- El archivo
tokenizer.model
estará en/LOCAL/PATH/TO/gemma/flax/2b-it/2
. - El punto de control del modelo estará en
/LOCAL/PATH/TO/gemma/flax/2b-it/2/2b-it
.
CKPT_PATH = os.path.join(GEMMA_PATH, GEMMA_VARIANT)
TOKENIZER_PATH = os.path.join(GEMMA_PATH, 'tokenizer.model')
print('CKPT_PATH:', CKPT_PATH)
print('TOKENIZER_PATH:', TOKENIZER_PATH)
CKPT_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/2b-it TOKENIZER_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/tokenizer.model
Realizar inferencias o muestreos
- Carga y formatea el punto de control del modelo Gemma con el método
gemma.params.load_and_format_params
:
from gemma import params as params_lib
params = params_lib.load_and_format_params(CKPT_PATH)
- Carga el tokenizador de Gemma, construido con
sentencepiece.SentencePieceProcessor
:
import sentencepiece as spm
vocab = spm.SentencePieceProcessor()
vocab.Load(TOKENIZER_PATH)
True
- Para cargar automáticamente la configuración correcta desde el punto de control del modelo de Gemma, usa
gemma.transformer.TransformerConfig
. El argumentocache_size
es la cantidad de pasos de tiempo en la cachéTransformer
de Gemma. Luego, crea una instancia del modelo de Gemma comotransformer
congemma.transformer.Transformer
(que se hereda deflax.linen.Module
).
from gemma import transformer as transformer_lib
transformer_config = transformer_lib.TransformerConfig.from_params(
params=params,
cache_size=1024
)
transformer = transformer_lib.Transformer(transformer_config)
- Crea un
sampler
congemma.sampler.Sampler
por encima de los puntos de control o los pesos del modelo de Gemma y el tokenizador:
from gemma import sampler as sampler_lib
sampler = sampler_lib.Sampler(
transformer=transformer,
vocab=vocab,
params=params['transformer'],
)
- Escribe una instrucción en
input_batch
y realiza inferencias. Puedes modificartotal_generation_steps
(la cantidad de pasos que se realizan cuando se genera una respuesta; en este ejemplo, se usa100
para preservar la memoria del host).
prompt = [
"\n# What is the meaning of life?",
]
reply = sampler(input_strings=prompt,
total_generation_steps=100,
)
for input_string, out_string in zip(prompt, reply.text):
print(f"Prompt:\n{input_string}\nOutput:\n{out_string}")
Prompt: # What is the meaning of life? Output: The question of what the meaning of life is one that has occupied the minds of philosophers, theologians, and individuals for centuries. There is no single, universally accepted answer, but there are many different perspectives on this complex and multifaceted question. **Some common perspectives on the meaning of life include:** * **Biological perspective:** From a biological standpoint, the meaning of life is to survive and reproduce. * **Existential perspective:** Existentialists believe that life is not inherently meaningful and that
- (Opcional) Ejecuta esta celda para liberar memoria si completaste el notebook y quieres probar con otra instrucción. Después, puedes volver a crear una instancia de
sampler
en el paso 3, y personalizar y ejecutar la instrucción en el paso 4.
del sampler
Más información
- Puedes obtener más información sobre la biblioteca
gemma
de Google DeepMind en GitHub, que contiene docstrings de módulos que usaste en este instructivo, comogemma.params
,gemma.transformer
ygemma.sampler
. - Las siguientes bibliotecas tienen sus propios sitios de documentación: core JAX, Flax y Orbax.
- Para ver la documentación de tokenizer/detokenizer de
sentencepiece
, consulta el repositorio de GitHub desentencepiece
de Google. - Para ver la documentación de
kagglehub
, consultaREADME.md
en el repositorio de GitHubkagglehub
de Kaggle. - Aprende a usar modelos de Gemma con Vertex AI de Google Cloud.