Ejecutar inferencias con Gemma mediante Keras

Ver en ai.google.dev Ejecutar en Google Colab Abrir en Vertex AI Ver el código fuente en GitHub

En este instructivo, se muestra cómo usar Gemma con KerasNLP para ejecutar inferencias y generar texto. Gemma es una familia de modelos abiertos ligeros y de última generación creados a partir de la misma investigación y tecnología que se usaron para crear los modelos de Gemini. KerasNLP es una colección de modelos de procesamiento de lenguaje natural (PLN) implementados en Keras que se pueden ejecutar en JAX, PyTorch y TensorFlow.

En este instructivo, usarás Gemma para generar respuestas de texto a varias instrucciones. Si es la primera vez que usas Keras, te recomendamos que leas Cómo comenzar a usar Keras antes de empezar, pero no es necesario. Aprenderás más sobre Keras a medida que trabajes en este instructivo.

Configuración

Configuración de Gemma

Para completar este instructivo, primero deberás completar las instrucciones de configuración de Gemma. Las instrucciones de configuración de Gemma te muestran cómo hacer lo siguiente:

  • Accede a Gemma en kaggle.com.
  • Selecciona un entorno de ejecución de Colab con recursos suficientes para ejecutar el modelo de Gemma 2B.
  • Generar y configurar un nombre de usuario Kaggle y una clave de API.

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.

Configura las variables de entorno

Configura variables de entorno para KAGGLE_USERNAME y KAGGLE_KEY.

import os
from google.colab import userdata

# Note: `userdata.get` is a Colab API. If you're not using Colab, set the env
# vars as appropriate for your system.
os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')

Instala dependencias

Instala Keras y KerasNLP.

# Install Keras 3 last. See https://keras.io/getting_started/ for more details.
pip install -q -U keras-nlp
pip install -q -U keras>=3

Seleccionar un backend

Keras es una API de aprendizaje profundo de varios marcos de alto nivel diseñada para ofrecer simplicidad y facilidad de uso. Keras 3 te permite elegir el backend: TensorFlow, JAX o PyTorch. Las tres opciones funcionarán en este instructivo.

import os

os.environ["KERAS_BACKEND"] = "jax"  # Or "tensorflow" or "torch".
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.9"

Importa paquetes

Importar Keras y KerasNLP.

import keras
import keras_nlp

Crea un modelo

KerasNLP proporciona implementaciones de muchas arquitecturas de modelos populares. En este instructivo, crearás un modelo con GemmaCausalLM, un modelo de Gemma de extremo a extremo para el modelado de lenguaje causal. Un modelo de lenguaje causal predice el siguiente token basándose en los anteriores.

Crea el modelo con el método from_preset:

gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma_2b_en")

Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...
Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...
Attaching 'model.weights.h5' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...
Attaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...
Attaching 'assets/tokenizer/vocabulary.spm' from model 'keras/gemma/keras/gemma_2b_en/1' to your Colab notebook...

La función GemmaCausalLM.from_preset() crea una instancia del modelo a partir de una arquitectura y pesos preestablecidos. En el código anterior, la cadena "gemma_2b_en" especifica el ajuste predeterminado del modelo Gemma 2B con 2,000 millones de parámetros. También están disponibles los modelos Gemma con parámetros 7B, 9B y 27B. Puedes encontrar las cadenas de código para los modelos de Gemma en sus fichas de Variación del modelo en kaggle.com.

Usa summary para obtener más información sobre el modelo:

gemma_lm.summary()

Como puedes ver en el resumen, el modelo tiene 2,500 millones de parámetros entrenables.

Generar texto

Ahora es el momento de generar algo de texto. El modelo tiene un método generate que genera texto basado en una instrucción. El argumento max_length opcional especifica la longitud máxima de la secuencia generada.

Pruébalo con el mensaje "What is the meaning of life?".

gemma_lm.generate("What is the meaning of life?", max_length=64)

'What is the meaning of life?\n\nThe question is one of the most important questions in the world.\n\nIt’s the question that has been asked by philosophers, theologians, and scientists for centuries.\n\nAnd it’s the question that has been asked by people who are looking for answers to their own lives'

Vuelve a llamar a generate con otra instrucción.

gemma_lm.generate("How does the brain work?", max_length=64)

'How does the brain work?\n\nThe brain is the most complex organ in the human body. It is responsible for controlling all of the body’s functions, including breathing, heart rate, digestion, and more. The brain is also responsible for thinking, feeling, and making decisions.\n\nThe brain is made up'

Si ejecutas en backends JAX o TensorFlow, notarás que la segunda llamada a generate se muestra casi al instante. Esto se debe a que cada llamada a generate para un tamaño de lote determinado y max_length se compila con XLA. La primera ejecución es costosa, pero las posteriores son mucho más rápidas.

También puedes proporcionar instrucciones por lotes usando una lista como entrada:

gemma_lm.generate(
    ["What is the meaning of life?",
     "How does the brain work?"],
    max_length=64)

['What is the meaning of life?\n\nThe question is one of the most important questions in the world.\n\nIt’s the question that has been asked by philosophers, theologians, and scientists for centuries.\n\nAnd it’s the question that has been asked by people who are looking for answers to their own lives',
 'How does the brain work?\n\nThe brain is the most complex organ in the human body. It is responsible for controlling all of the body’s functions, including breathing, heart rate, digestion, and more. The brain is also responsible for thinking, feeling, and making decisions.\n\nThe brain is made up']

Opcional: Prueba con otra muestra

Puedes controlar la estrategia de generación de GemmaCausalLM si estableces el argumento sampler en compile(). De forma predeterminada, se usará el muestreo "greedy".

Como experimento, intenta establecer una estrategia "top_k":

gemma_lm.compile(sampler="top_k")
gemma_lm.generate("What is the meaning of life?", max_length=64)

'What is the meaning of life? That was a question I asked myself as I was driving home from work one night in 2012. I was driving through the city of San Bernardino, and all I could think was, “What the heck am I doing?”\n\nMy life was completely different. I'

Mientras que el algoritmo codicioso predeterminado siempre elige el token con la mayor probabilidad, el algoritmo Top-K elige de forma aleatoria el siguiente token entre los tokens de probabilidad de Top-K.

No es necesario que especifiques una muestra y puedes ignorar el último fragmento de código si no es útil para tu caso de uso. Si deseas obtener más información sobre los muestras disponibles, consulta Samplers.

¿Qué sigue?

En este instructivo, aprendiste a generar texto con KerasNLP y Gemma. Estas son algunas sugerencias sobre lo que puedes aprender a continuación: