Guia de início rápido do Keras CodeGemma

Ver em ai.google.dev Executar no Google Colab Veja o código-fonte no GitHub

O CodeGemma é uma família de modelos abertos, leves e modernos, criados com a mesma pesquisa e tecnologia usada para criar os modelos do Gemini.

Os modelos do CodeGemma são treinados com mais de 500 bilhões de tokens, principalmente de código, usando as mesmas arquiteturas da família de modelos Gemma. Como resultado, os modelos do CodeGemma alcançam um desempenho de código de última geração em ambos e geração de cargas de trabalho, mantendo uma sólida compreensão e raciocínio em escala.

O CodeGemma tem três variantes:

  • Um modelo pré-treinado com códigos 7B
  • Um modelo de código ajustado por instruções 7B
  • Um modelo 2B, treinado especificamente para preenchimento de código e geração aberta.

Neste guia, você vai aprender a usar o modelo 2B do CodeGemma com KerasNLP para uma tarefa de preenchimento de código.

Configuração

Acessar o CodeGemma

Para concluir este tutorial, primeiro você precisa concluir as instruções de configuração na 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 ser executado 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, vá para a próxima seção, em que você definirá variáveis de ambiente para o ambiente do Colab.

Selecione o ambiente de execução

Para concluir este tutorial, você precisa de um ambiente de execução do Colab com recursos suficientes para executar o modelo CodeGemma 2B. Nesse caso, você pode usar uma GPU T4:

  1. No canto superior direito da janela do Colab, selecione ▾ (Opções de conexão adicionais).
  2. Selecione Mudar o tipo de ambiente de execução.
  3. Em Acelerador de hardware, selecione GPU T4.

Configurar sua chave de API

Para usar o Gemma, forneça seu nome de usuário do Kaggle e uma chave de API do Kaggle.

Para gerar uma chave de API do Kaggle, vá para a guia Conta do seu perfil de usuário do Kaggle e selecione Criar novo token. Isso acionará o download de um arquivo kaggle.json com suas credenciais de API.

No Colab, selecione Secrets (Eu) no painel esquerdo e adicione seu nome de usuário e chave de API do Kaggle. Armazene seu nome de usuário com o nome KAGGLE_USERNAME e sua chave de API com o nome KAGGLE_KEY.

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

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

Instalar dependências

pip install -q -U keras-nlp

Selecione um back-end

Keras é uma API de aprendizado profundo de alto nível e com vários frameworks projetadas para simplicidade e facilidade de uso. Com o Keras 3, é possível executar fluxos de trabalho em um destes três back-ends: TensorFlow, JAX ou PyTorch.

Para este tutorial, configure o back-end do TensorFlow.

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

Importar pacotes

Importe o Keras e o KerasNLP.

import keras_nlp
import keras

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

Carregar modelo

O KerasNLP oferece implementações de várias 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("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]

O método from_preset instancia o modelo usando uma arquitetura e pesos predefinidos. No código acima, a string code_gemma_2b_en especifica a arquitetura predefinida: um modelo CodeGemma com 2 bilhões de parâmetros.

Preenchimento automático de código preenchimento no meio

Este exemplo usa o recurso de preenchimento no meio (FIM, na sigla em inglês) do CodeGemma para concluir o código com base no contexto. Isso é particularmente útil em aplicativos de edição de código para inserir código onde o cursor de texto é baseado no código ao redor dele (antes e depois do cursor).

O CodeGemma permite usar quatro tokens definidos pelo usuário: três para FIM e um token <|file_separator|> para suporte a contexto de vários arquivos. Use-os para definir constantes.

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

Defina os tokens de parada para o 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)

Formate o comando para preenchimento de código. Observações:

  • Não deve haver espaços em branco entre os tokens do FIM e o prefixo e sufixo
  • O token do meio do FIM precisa estar no final para preparar o modelo e continuar preenchendo
  • O prefixo ou sufixo pode estar vazio dependendo de onde o cursor está no arquivo no momento ou de quanto contexto você quer dar ao modelo.

Use uma função auxiliar para formatar o comando.

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

Execute o comando. É recomendável transmitir tokens de resposta. Interrompa o streaming ao encontrar qualquer um dos tokens definidos pelo usuário ou de fim de turno/senetence para conseguir o preenchimento de 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|>'

O modelo fornece sys como o preenchimento de código sugerido.

Resumo

Este tutorial mostrou como usar o CodeGemma para preencher o código com base no contexto. Em seguida, confira o notebook de programação assistida por IA com o CodeGemma e o KerasNLP (em inglês) para mais exemplos de como usar o CodeGemma.

Consulte também o card de modelo do CodeGemma para conferir as especificações técnicas dos modelos do CodeGemma.