查看 Gemma Cookbook 代码库，了解生成和优化示例！了解详情

此页面由 Cloud Translation API 翻译。

Keras CodeGemma 快速入门

前往 ai.google.dev 查看

在 Google Colab 中运行

在 GitHub 上查看源代码

CodeGemma 是一系列先进的轻量级开放模型，采用与 Gemini 模型相同的研究和技术构建而成。

CodeGemma 模型使用 5, 000 多亿个主要是代码词元的词元进行训练，与 Gemma 模型系列相同的架构。因此，CodeGemma 模型在生成任务和生成任务，理解能力和推理能力。

CodeGemma 有 3 个变体：

70 亿个代码的预训练模型
70 亿指令调优的代码模型
一个 2B 模型，专为代码填充和开放式生成而训练。

本指南将逐步介绍如何将 CodeGemma 2B 模型与 KerasNLP 结合使用来完成代码补全任务。

设置

访问 CodeGemma

要完成本教程，您首先需要在 Gemma 设置中完成设置说明。Gemma 设置说明介绍了如何执行以下操作：

在 kaggle.com 上访问 Gemma。
选择具有足够资源的 Colab 运行时来运行 Gemma 2B 模型。
生成并配置 Kaggle 用户名和 API 密钥。

完成 Gemma 设置后，请继续执行下一部分，您将为 Colab 环境设置环境变量。

选择运行时

要完成本教程，您需要有一个具有足够资源的 Colab 运行时来运行 CodeGemma 2B 模型。在这种情况下，您可以使用 T4 GPU：

在 Colab 窗口的右上角，选择 ▾（其他连接选项）。
选择更改运行时类型。
在硬件加速器下，选择 T4 GPU。

配置您的 API 密钥

要使用 Gemma，您必须提供您的 Kaggle 用户名和 Kaggle API 密钥。

要生成 Kaggle API 密钥，请前往您的 Kaggle 用户个人资料中的 Account（账号）标签页，然后选择 Create New Token（创建新令牌）。这会触发下载包含 API 凭据的 kaggle.json 文件。

在 Colab 中，选择左侧窗格中的 Secrets (🔑?)，然后添加您的 Kaggle 用户名和 Kaggle API 密钥。将您的用户名存储在名称 KAGGLE_USERNAME 下，将 API 密钥存储在名称 KAGGLE_KEY 下。

设置环境变量

为 KAGGLE_USERNAME 和 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')

安装依赖项

pip install -q -U keras-nlp

选择一个后端

Keras 是一个高级的多框架深度学习 API，旨在实现简洁易用。使用 Keras 3，您可以在以下三种后端之一上运行工作流：TensorFlow、JAX 或 PyTorch。

在本教程中，请为 TensorFlow 配置后端。

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

导入软件包

导入 Keras 和 KerasNLP。

import keras_nlp
import keras

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

加载模型

KerasNLP 提供了许多热门模型架构的实现。在本教程中，您将使用 GemmaCausalLM 创建一个模型，这是一个用于因果语言建模的端到端 Gemma 模型。因果语言模型会根据上一个词元预测下一个词元。

使用 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]

from_preset 方法会根据预设架构和权重对模型进行实例化。在上面的代码中，字符串 code_gemma_2b_en 指定了预设架构，即一个具有 20 亿参数的 CodeGemma 模型。

填充中间代码补全

此示例使用 CodeGemma 的中间填充 (FIM) 功能根据周围的上下文完成代码。这在代码编辑器应用中插入代码时特别有用，而且这些代码会在文本光标位于其周围代码的位置（光标前后）显示。

借助 CodeGemma，您可以使用 4 个用户定义的令牌：3 个用于 FIM 令牌，1 个 <|file_separator|> 令牌用于多文件上下文支持。使用这些变量定义常量。

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

为模型定义停止令牌。

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)

代码补全提示的格式。请注意：

任何 FIM 令牌与前缀和后缀之间不应有空格
FIM 中间令牌应位于末尾，以启动模型以继续填充
前缀或后缀可以为空，具体取决于游标目前在文件中的位置，或者您想要为模型提供多少上下文

使用辅助函数设置提示的格式。

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

运行提示。建议流式传输响应令牌。遇到任何用户定义令牌或回合结束令牌/序列令牌时，停止流式传输，以获得完整代码。

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

该模型提供 sys 作为建议的代码补全功能。

摘要

本教程为您介绍了如何使用 CodeGemma 根据周围的上下文填充代码。接下来，请查看《使用 CodeGemma 和 KerasNLP 进行 AI 辅助编程》(AI Assisted Programming with CodeGemma) 和 KerasNLP 笔记本，了解更多有关如何使用 CodeGemma 的示例。

如需了解 CodeGemma 模型的技术规范，另请参阅 CodeGemma 模型卡片。