Começar a usar o Gemma com o KerasNLP

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

Este tutorial mostra como começar a usar o Gemma com o KerasNLP. O Gemma é uma família de modelos abertos leves e de última geração criados com a mesma pesquisa e tecnologia usadas 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 a várias solicitações. Se você não conhece o Keras, leia Primeiros passos com o Keras antes de começar, mas não é obrigatório. Você vai aprender mais sobre o Keras ao longo deste tutorial.

Configuração

Configuração do Gemma

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

  • Acesse o 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ê vai definir variáveis de ambiente para o 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

Instale 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"

Selecione um back-end

O Keras é uma API de aprendizado profundo de alto nível e com vários frameworks projetada para ser simples e fácil de usar. O Keras 3 permite escolher o back-end: TensorFlow, JAX ou PyTorch. Todos os três funcionam para este tutorial.

import os

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

Importar pacotes

Importe o Keras e o KerasNLP.

import keras
import keras_nlp

Criar um modelo

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

Crie o modelo usando o método from_preset:

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

A função GemmaCausalLM.from_preset() instancia o modelo usando uma arquitetura e pesos predefinidos. No código acima, a string "gemma2_2b_en" especifica a predefinição do modelo Gemma 2 2B com 2 bilhões de parâmetros. Os modelos Gemma com parâmetros 7B, 9B e 27B também estão disponíveis. As strings de código dos modelos Gemma estão nas listagens de variação de modelo no Kaggle.

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

gemma_lm.summary()

Como você pode conferir no resumo, o modelo tem 2,6 bilhões de parâmetros treináveis.

Gerar texto

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

Teste com o comando "what is keras in 3 bullet points?".

gemma_lm.generate("what is keras in 3 bullet points?", max_length=64)

'what is keras in 3 bullet points?\n\n[Answer 1]\n\nKeras is a high-level neural networks API, written in Python and capable of running on top of TensorFlow, CNTK, Theano, or PlaidML. It is designed to be user-friendly and easy to extend.\n\n'

Tente chamar generate novamente com um comando diferente.

gemma_lm.generate("The universe is", max_length=64)

'The universe is a vast and mysterious place, filled with countless stars, planets, and galaxies. But what if there was a way to see the universe in a whole new way? What if we could see the universe as it was when it was first created? What if we could see the universe as it is now'

Se você estiver executando em back-ends JAX ou TensorFlow, vai notar que a segunda chamada generate retorna 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 subsequentes são muito mais rápidas.

Também é possível fornecer comandos em lote usando uma lista como entrada:

gemma_lm.generate(
    ["what is keras in 3 bullet points?",
     "The universe is"],
    max_length=64)

['what is keras in 3 bullet points?\n\n[Answer 1]\n\nKeras is a high-level neural networks API, written in Python and capable of running on top of TensorFlow, CNTK, Theano, or PlaidML. It is designed to be user-friendly and easy to extend.\n\n',
 'The universe is a vast and mysterious place, filled with countless stars, planets, and galaxies. But what if there was a way to see the universe in a whole new way? What if we could see the universe as it was when it was first created? What if we could see the universe as it is now']

Opcional: testar um sampler diferente

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

Como um experimento, tente definir uma estratégia "top_k":

gemma_lm.compile(sampler="top_k")
gemma_lm.generate("The universe is", max_length=64)

'The universe is a big place, and there are so many things we do not know or understand about it.\n\nBut we can learn a lot about our world by studying what is known to us.\n\nFor example, if you look at the moon, it has many features that can be seen from the surface.'

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

Não é necessário especificar um sampler, e você pode ignorar 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 o KerasNLP e o Gemma. Confira algumas sugestões do que aprender a seguir: