Programação assistida por IA com CodeGemma e KerasNLP

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

Visão geral

O CodeGemma é uma variante do Gemma ajustada para tarefas de programação. Este tutorial baseia-se no guia de início rápido do Keras CodeGemma (link em inglês) e mostra mais maneiras como o CodeGemma pode ajudar nas tarefas de programação.

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 executar o modelo Gemma 7B.
  • 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 executar os modelos CodeGemma 7B, você precisa ter um plano Colab Pro pago que ofereça um ambiente de execução com uma GPU A100.

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

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 para JAX.

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

Importar pacotes

Importe o Keras e o KerasNLP.

import keras_nlp
import keras

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

Exemplos de modelo do CodeGemma 7B

Nesta seção, abordamos exemplos de uso do modelo 7B CodeGemma pré-treinado para ajudar em tarefas de programação.

Carregar o modelo

O KerasNLP fornece implementações das três variantes do CodeGemma (2B e 7B pré-treinadas (PT, na sigla em inglês) e 7B ajustadas por instruções (IT)) 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.

Neste exemplo, carregue o modelo code_gemma_7b_en usando o método from_preset.

gemma_lm_7b = keras_nlp.models.GemmaCausalLM.from_preset("code_gemma_7b_en")
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/config.json...
100%|██████████| 556/556 [00:00<00:00, 790kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/model.weights.h5...
100%|██████████| 15.9G/15.9G [02:39<00:00, 107MB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/tokenizer.json...
100%|██████████| 401/401 [00:00<00:00, 587kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/assets/tokenizer/vocabulary.spm...
100%|██████████| 4.04M/4.04M [00:00<00:00, 16.4MB/s]
gemma_lm_7b.summary()

O método from_preset instancia o modelo usando uma arquitetura e pesos predefinidos.

Preenchimento automático de código com FIM de várias linhas

Os modelos PT CodeGemma são treinados em tarefas de preenchimento de código. Esta seção mostra exemplos que usam o recurso de preenchimento no meio de várias linhas (FIM, na sigla em inglês) do CodeGemma para preencher automaticamente o código no local especificado do cursor com base no contexto.

Como primeira etapa, defina as constantes e uma função auxiliar de formatação do prompt.

# Formatting control tokens to specify cursor location
BEFORE_CURSOR = "<|fim_prefix|>"
AFTER_CURSOR = "<|fim_suffix|>"
AT_CURSOR = "<|fim_middle|>"
FILE_SEPARATOR = "<|file_separator|>"

# Define model stop tokens
END_TOKEN = gemma_lm_7b.preprocessor.tokenizer.end_token
stop_tokens = (BEFORE_CURSOR, AFTER_CURSOR, AT_CURSOR, FILE_SEPARATOR, END_TOKEN)
stop_token_ids = tuple(gemma_lm_7b.preprocessor.tokenizer.token_to_id(x) for x in stop_tokens)

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

Exemplo 1: inserir uma condição ausente

O código de exemplo abaixo para gerar a sequência de Fibonacci não será executado corretamente se n=1:

def fibonacci(n: int) -> int:
  if n == 0:
    return 0
  # The cursor is right before the e in the following line
  else:
    return fibonacci(n - 1) + fibonacci(n - 2)

Supondo que o cursor esteja no início da linha 4 (em que a cláusula else está), o conteúdo antes e depois dele será:

before = """def fibonacci(n: int) -> int:\n  if n == 0:\n    return 0\n""" # Mind the spaces!
after = """\n  else:\n    return fibonacci(n - 1) + fibonacci(n-2)\n"""
prompt = format_completion_prompt(before, after)
print(prompt)
<|fim_prefix|>def fibonacci(n: int) -> int:
  if n == 0:
    return 0
<|fim_suffix|>
  else:
    return fibonacci(n - 1) + fibonacci(n-2)
<|fim_middle|>

Execute o comando.

print(gemma_lm_7b.generate(prompt, stop_token_ids=stop_token_ids, max_length=128))
<|fim_prefix|>def fibonacci(n: int) -> int:
  if n == 0:
    return 0
<|fim_suffix|>
  else:
    return fibonacci(n - 1) + fibonacci(n-2)
<|fim_middle|>elif n == 1:
    return 1<|file_separator|>

O modelo insere a conidação elif correta para n=1 no local do cursor.

Exemplo 2: algoritmo completo de travessia de DFS

Código de preenchimento automático para um algoritmo de travessia em árvore de pesquisa em profundidade (DFS, na sigla em inglês).

before = """void dfs(node* root) {
  if (root->left) {
    dfs(root->left);
  }"""
after = """\nprintf("%d", root->value);
}"""
prompt = format_completion_prompt(before, after)
print(prompt)
<|fim_prefix|>void dfs(node* root) {
  if (root->left) {
    dfs(root->left);
  }<|fim_suffix|>
printf("%d", root->value);
}<|fim_middle|>

Execute o comando.

print(gemma_lm_7b.generate(prompt, stop_token_ids=stop_token_ids, max_length=128))
<|fim_prefix|>void dfs(node* root) {
  if (root->left) {
    dfs(root->left);
  }<|fim_suffix|>
printf("%d", root->value);
}<|fim_middle|>
  if (root->right) {
    dfs(root->right);
  }<|file_separator|>

Geração de códigos

Além do preenchimento de código, o CodeGemma 7B PT também é treinado em corpus de linguagem natural. É possível usar isso para solicitar que o modelo gere código.

generation_prompt= """Write a rust function to identify non-prime numbers.
Examples:
>>> is_not_prime(2)
False
>>> is_not_prime(10)
True
pub fn is_not_prime(n: i32) -> bool {"""
print(gemma_lm_7b.generate(generation_prompt, max_length=500))
Write a rust function to identify non-prime numbers.
Examples:
>>> is_not_prime(2)
False
>>> is_not_prime(10)
True
pub fn is_not_prime(n: i32) -> bool {
    if n <= 1 {
        return true;
    }
    for i in 2..n {
        if n % i == 0 {
            return true;
        }
    }
    false
}

Exemplos de modelos de TI 7B

Esta seção usa o modelo CodeGemma 7B com orientação para tarefas de programação mais avançadas. O modelo de TI 7B do CodeGemma é derivado do modelo PT do CodeGemma 7B por meio de ajuste fino supervisionado no código junto com o aprendizado por reforço com feedback humano. Nesta seção, abordamos exemplos de uso desse modelo para geração aberta.

Carregue o modelo de TI

Carregue o modelo code_gemma_instruct_7b_en usando o método from_preset.

gemma_lm_7b_it = keras_nlp.models.GemmaCausalLM.from_preset("code_gemma_instruct_7b_en")
gemma_lm_7b_it.summary()
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/config.json...
100%|██████████| 556/556 [00:00<00:00, 754kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/model.weights.h5...
100%|██████████| 15.9G/15.9G [03:18<00:00, 86.2MB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/tokenizer.json...
100%|██████████| 401/401 [00:00<00:00, 593kB/s]
Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/assets/tokenizer/vocabulary.spm...
100%|██████████| 4.04M/4.04M [00:00<00:00, 16.8MB/s]

Os modelos de TI são treinados com um formatador específico que anota todos os exemplos de ajuste de instruções com informações extras para indicar funções e delinear os turnos em uma conversa.

Como primeira etapa, defina as constantes e uma função auxiliar de formatação do prompt.

# Formatting control tokens for instruction tuning
START_OF_TURN_USER = "<start_of_turn>user"
END_OF_TURN = "<end_of_turn>"
START_OF_TURN_MODEL = "<start_of_turn>model"

# Formatting helper function
def format_instruction_prompt(context):
    return f"{START_OF_TURN_USER}\n{context}{END_OF_TURN}\n{START_OF_TURN_MODEL}\n"

Tradução de código

context1 = """
You are an experienced C and Python programmer. Convert the following Python code into C.
```python
def factorial(n):
    result = 1
    for i in range(2, n + 1):
        result *= i
    return result
```\n"""

Formate o comando.

prompt1 = format_instruction_prompt(context1)
print(prompt1)
<start_of_turn>user

You are an experienced C and Python programmer. Convert the following Python code into C.

```python
def factorial(n):
    result = 1
    for i in range(2, n + 1):
        result *= i
    return result
```
<end_of_turn>
<start_of_turn>model

Execute o comando.

print(gemma_lm_7b_it.generate(prompt1, max_length=500))
<start_of_turn>user

You are an experienced C and Python programmer. Convert the following Python code into C.

```python
def factorial(n):
    result = 1
    for i in range(2, n + 1):
        result *= i
    return result
```
<end_of_turn>
<start_of_turn>model
Here is the C code equivalent of the Python code:

```c
int factorial(int n) {
  int result = 1;
  for (int i = 2; i <= n; i++) {
    result *= i;
  }
  return result;
}
```

Here is a breakdown of the changes:

* The function is declared with the `int` return type, as in Python.
* The `for` loop is converted to a `for` loop with an `int` variable `i` initialized to 2 and incremented by 1 in each iteration.
* The `range` function is replaced with a simple loop that iterates from 2 to `n` (inclusive).
* The `result *= i` statement is used to multiply `result` by `i` in each iteration.
* The `return` statement is used to return the final value of `result`.

Detecção de vulnerabilidade de código

context2 = """
You are an experienced C++ programmer hunting for vulnerable code. Is the following code vulnerable? Explain your reasoning.
```cpp
int i;
unsigned int numWidgets;
Widget **WidgetList;

numWidgets = GetUntrustedSizeValue();
if ((numWidgets == 0) || (numWidgets > MAX_NUM_WIDGETS)) {
    ExitError("Incorrect number of widgets requested!");
}
WidgetList = (Widget **) malloc(numWidgets * sizeof(Widget *));
printf("WidgetList ptr=%p\n", WidgetList);
for (i = 0; i < numWidgets; i++) {
    WidgetList[i] = InitializeWidget();
}
WidgetList[numWidgets] = NULL;
showWidgets(WidgetList);
```\n"""

Formate o comando.

prompt2 = format_instruction_prompt(context2)
print(prompt2)
<start_of_turn>user

You are an experienced C++ programmer hunting for vulnerable code. Is the following code vulnerable? Explain your reasoning.

```cpp
int i;
unsigned int numWidgets;
Widget **WidgetList;

numWidgets = GetUntrustedSizeValue();
if ((numWidgets == 0) || (numWidgets > MAX_NUM_WIDGETS)) {
    ExitError("Incorrect number of widgets requested!");
}
WidgetList = (Widget **) malloc(numWidgets * sizeof(Widget *));
printf("WidgetList ptr=%p
", WidgetList);
for (i = 0; i < numWidgets; i++) {
    WidgetList[i] = InitializeWidget();
}
WidgetList[numWidgets] = NULL;
showWidgets(WidgetList);
```
<end_of_turn>
<start_of_turn>model
print(gemma_lm_7b_it.generate(prompt2, max_length=1000))
<start_of_turn>user

You are an experienced C++ programmer hunting for vulnerable code. Is the following code vulnerable? Explain your reasoning.

```cpp
int i;
unsigned int numWidgets;
Widget **WidgetList;

numWidgets = GetUntrustedSizeValue();
if ((numWidgets == 0) || (numWidgets > MAX_NUM_WIDGETS)) {
    ExitError("Incorrect number of widgets requested!");
}
WidgetList = (Widget **) malloc(numWidgets * sizeof(Widget *));
printf("WidgetList ptr=%p
", WidgetList);
for (i = 0; i < numWidgets; i++) {
    WidgetList[i] = InitializeWidget();
}
WidgetList[numWidgets] = NULL;
showWidgets(WidgetList);
```
<end_of_turn>
<start_of_turn>model
Yes, the code is vulnerable to a memory access error.

**Reasoning:**

* The code allocates memory for `WidgetList` using `malloc` based on the value of `numWidgets`.
* However, the loop iterates from `0` to `numWidgets`, which is one element beyond the allocated memory.
* This means that accessing `WidgetList[numWidgets]` will result in a memory access error, as it is outside the bounds of the allocated memory.

**Example of Memory Access Error:**

When `numWidgets` is 5, the code allocates memory for `WidgetList` as follows:

```
WidgetList = (Widget **) malloc(5 * sizeof(Widget *));
```

The loop iterates from 0 to 4, accessing the following elements:

* `WidgetList[0]`
* `WidgetList[1]`
* `WidgetList[2]`
* `WidgetList[3]`
* `WidgetList[4]`

However, the code then attempts to access `WidgetList[5]`, which is outside the allocated memory range. This will result in a memory access error.

**Solution:**

To resolve this vulnerability, the loop should be modified to iterate from 0 to `numWidgets - 1`:

```cpp
for (i = 0; i < numWidgets - 1; i++) {
    WidgetList[i] = InitializeWidget();
}
```

This ensures that the loop does not access elements beyond the allocated memory range.

O modelo detecta uma possível vulnerabilidade no código e faz mudanças no código para mitigá-la.

Resumo

Este tutorial mostrou como usar o CodeGemma em diversas tarefas de programação. Para saber mais sobre o CodeGemma: