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:

  1. 選取 Colab 視窗右上方的「▾」(其他連線選項)
  2. 選取「變更執行階段類型」
  3. 在「硬體加速器」下方,選取「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_USERNAMEKAGGLE_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 的權杖),以及一個 <|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 輔助程式設計」一文,取得更多使用 CodeGemma 的範例。

如要瞭解 CodeGemma 模型的技術規格,請參閱 CodeGemma 模型卡