Guía de inicio rápido de Keras CodeGemma

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

CodeGemma 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.

Los modelos de CodeGemma se entrenan con más de 500,000 millones de tokens de código principalmente, usando con las mismas arquitecturas que la familia de modelos Gemma. Como resultado, los modelos de CodeGemma logran un rendimiento de código de vanguardia y generación de proyectos, mientras se mantienen habilidades de comprensión y razonamiento a gran escala.

CodeGemma tiene 3 variantes:

  • Un modelo previamente entrenado en código 7B
  • Un modelo 7B de código ajustado a instrucciones
  • Un modelo 2B, entrenado específicamente para el relleno de código y la generación abierta.

En esta guía, se explica cómo usar el modelo 2B de CodeGemma con KerasNLP para realizar una tarea de finalización de código.

Configuración

Accede a CodeGemma

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:

  • Obtén acceso a Gemma en kaggle.com.
  • Selecciona un entorno de ejecución de Colab con recursos suficientes para ejecutarlo el modelo 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.

Selecciona el entorno de ejecución

Para completar este tutorial, necesitarás un entorno de ejecución de Colab con recursos suficientes para ejecutar el modelo 2B de CodeGemma. En este caso, puedes usar una GPU T4:

  1. En la esquina superior derecha de la ventana de Colab, selecciona ▾ (Opciones de conexión adicionales).
  2. Selecciona Cambiar tipo de entorno de ejecución.
  3. En Acelerador de hardware, selecciona GPU T4.

Cómo configurar tu clave de API

Para usar Gemma, debes proporcionar tu nombre de usuario y una clave de API de Kaggle.

Para generar una clave de API de Kaggle, dirígete a la pestaña Cuenta de tu perfil de usuario de Kaggle y selecciona Crear token nuevo. Esto activará la descarga de un archivo kaggle.json que contiene tus credenciales de API.

En Colab, selecciona Secrets (Compose) en el panel izquierdo y agrega tu nombre de usuario de Kaggle y tu clave de API de Kaggle. Almacena tu nombre de usuario con el nombre KAGGLE_USERNAME y tu clave de API con el nombre KAGGLE_KEY.

Configure las variables de entorno

Configura variables de entorno para KAGGLE_USERNAME y KAGGLE_KEY.

import os
from google.colab import userdata

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

Instala dependencias

pip install -q -U keras-nlp

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. Con Keras 3, puedes ejecutar flujos de trabajo en uno de los tres backends: TensorFlow, JAX o PyTorch.

En este instructivo, configura el backend para TensorFlow.

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

Importa paquetes

Importar Keras y KerasNLP.

import keras_nlp
import keras

# Run at half precision.
keras.config.set_floatx("bfloat16")

Cargar 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("code_gemma_2b_en")
gemma_lm.summary()
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_2b_en/1/download/config.json...
100%|██████████| 554/554 [00:00<00:00, 1.41MB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_2b_en/1/download/model.weights.h5...
100%|██████████| 4.67G/4.67G [05:06<00:00, 16.4MB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_2b_en/1/download/tokenizer.json...
100%|██████████| 401/401 [00:00<00:00, 382kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_2b_en/1/download/assets/tokenizer/vocabulary.spm...
100%|██████████| 4.04M/4.04M [00:01<00:00, 2.41MB/s]

El método from_preset crea una instancia del modelo a partir de una arquitectura y pesos preestablecidos. En el código anterior, la cadena code_gemma_2b_en especifica la arquitectura predeterminada: un modelo de CodeGemma con 2,000 millones de parámetros.

Finalización de código de completado de intermedio

Este ejemplo usa la capacidad de completar el medio (FIM) de CodeGemma para completar el código según el contexto circundante. Esto es particularmente útil en aplicaciones de editor de código para insertar código donde el cursor de texto se basa en el código que lo rodea (antes y después del cursor).

CodeGemma te permite usar 4 tokens definidos por el usuario: 3 para FIM y un token <|file_separator|> para admitir el contexto de varios archivos. Úsalas para definir constantes.

BEFORE_CURSOR = "<|fim_prefix|>"
AFTER_CURSOR = "<|fim_suffix|>"
AT_CURSOR = "<|fim_middle|>"
FILE_SEPARATOR = "<|file_separator|>"

Define los tokens de detención para el modelo.

END_TOKEN = gemma_lm.preprocessor.tokenizer.end_token

stop_tokens = (BEFORE_CURSOR, AFTER_CURSOR, AT_CURSOR, FILE_SEPARATOR, END_TOKEN)

stop_token_ids = tuple(gemma_lm.preprocessor.tokenizer.token_to_id(x) for x in stop_tokens)

Dar formato al mensaje para completar el código Ten en cuenta lo siguiente:

  • No debe haber espacios en blanco entre los tokens de FIM y el prefijo y el sufijo
  • El token del medio de FIM debe estar al final para preparar el modelo y seguir completando
  • El prefijo o sufijo puede estar vacío según la ubicación del cursor en el archivo o la cantidad de contexto que quieras darle al modelo.

Usa una función auxiliar para darle formato a la instrucción.

def format_completion_prompt(before, after):
    return f"{BEFORE_CURSOR}{before}{AFTER_CURSOR}{after}{AT_CURSOR}"

before = "import "
after = """if __name__ == "__main__":\n    sys.exit(0)"""
prompt = format_completion_prompt(before, after)
print(prompt)
<|fim_prefix|>import <|fim_suffix|>if __name__ == "__main__":
    sys.exit(0)<|fim_middle|>

Ejecuta la instrucción. Se recomienda transmitir tokens de respuesta. Detén la transmisión cuando encuentres uno de los tokens definidos por el usuario, o de fin de turno o secuencia, para completar el código resultante.

gemma_lm.generate(prompt, stop_token_ids=stop_token_ids, max_length=128)
'<|fim_prefix|>import <|fim_suffix|>if __name__ == "__main__":\n    sys.exit(0)<|fim_middle|>sys\n<|file_separator|>'

El modelo proporciona sys como la finalización de código sugerida.

Resumen

Este tutorial te explicó cómo utilizar CodeGemma para rellenar código según el contexto circundante. A continuación, consulta el notebook de programación asistida por IA con CodeGemma y KerasNLP para obtener más ejemplos sobre cómo puedes usar CodeGemma.

Consulta también la tarjeta de modelo de CodeGemma para conocer las especificaciones técnicas de los modelos de CodeGemma.