Começar a usar o Gemma com o KerasNLP

Ver em ai.google.dev Executar no Google Colab Abrir na Vertex AI Consulte o código-fonte no GitHub

Neste tutorial, mostramos como começar a usar o Gemma usando o KerasNLP. A Gemma é uma família de modelos abertos leves e modernos, criados a partir da mesma pesquisa e tecnologia usada na criação dos modelos do Gemini. O KerasNLP é uma coleção de modelos de processamento de linguagem natural (PLN) implementados no Keras e executáveis no JAX, PyTorch e TensorFlow.

Neste tutorial, você vai usar o Gemma para gerar respostas de texto para vários comandos. Se você é novo no Keras, talvez queira ler Getting started with Keras antes de começar, mas não é obrigatório. Você aprenderá mais sobre o Keras durante este tutorial.

Configuração

Configuração do Gemma

Para concluir este tutorial, primeiro você precisa concluir as instruções de configuração em Gemma setup. As instruções de configuração do Gemma mostram como fazer o seguinte:

  • Obtenha acesso à Gemma em kaggle.com.
  • Selecione um ambiente de execução do Colab com recursos suficientes para executar o modelo Gemma 2B.
  • Gere e configure um nome de usuário e uma chave de API do Kaggle.

Depois de concluir a configuração do Gemma, passe para a próxima seção, em que você definirá variáveis para seu ambiente do Colab.

Defina as variáveis de ambiente

Defina as variáveis de ambiente para KAGGLE_USERNAME e 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')

Instalar dependências

Instalar o Keras e o 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

Selecionar um back-end

A Keras é uma API de aprendizado profundo de alto nível com vários frameworks, criada para simplificar e facilitar o uso. O Keras 3 permite escolher o back-end: TensorFlow, JAX ou PyTorch. Os três funcionarão neste tutorial.

import os

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

Importar pacotes

Importar Keras e KerasNLP.

import keras
import keras_nlp

Criar um modelo

O KerasNLP oferece implementações de muitas arquiteturas de modelo conhecidas. Neste tutorial, você criará um modelo usando GemmaCausalLM, um modelo Gemma completo para modelagem de linguagem causal. Um modelo de linguagem causal prevê o próximo token com base em tokens anteriores.

Crie o modelo usando o 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...

from_preset instancia o modelo com base em uma arquitetura e pesos predefinidos. No código acima, a string "gemma_2b_en" especifica a arquitetura predefinida: um modelo Gemma com dois bilhões de parâmetros.

Use summary para ter mais informações sobre o modelo:

gemma_lm.summary()

Como é possível notar no resumo, o modelo tem 2,5 bilhões de parâmetros treináveis.

Gerar texto

Agora é hora de gerar texto. O modelo tem um método generate que gera texto com base em um comando. O argumento max_length opcional especifica o tamanho máximo da sequência gerada.

Teste com o comando "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'

Tente chamar generate novamente com outro comando.

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'

Se você estiver executando em back-ends JAX ou TensorFlow, vai notar que a segunda chamada generate é retornada quase instantaneamente. Isso ocorre porque cada chamada para generate para um determinado tamanho de lote e max_length é compilada com XLA. A primeira execução é cara, mas as execuções subsequentes são muito mais rápidas.

Também é possível enviar comandos em lote usando uma 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: tente usar outro sampler

É possível controlar a estratégia de geração para GemmaCausalLM definindo o argumento sampler em compile(). Por padrão, será usada a amostragem de "greedy".

Como experimento, defina uma estratégia "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'

Enquanto o algoritmo ganancioso padrão sempre escolhe o token com a maior probabilidade, o algoritmo top-K seleciona aleatoriamente o próximo token entre os tokens de probabilidade top-k.

Não é necessário especificar um sampler. Ignore o último snippet de código se ele não for útil para seu caso de uso. Para saber mais sobre os samplers disponíveis, consulte Samplers.

A seguir

Neste tutorial, você aprendeu a gerar texto usando KerasNLP e Gemma. Aqui estão algumas sugestões do que você pode aprender a seguir: