![]() | ![]() | ![]() |
Мы представляем CodeGemma, коллекцию моделей открытого кода, основанную на моделях Gemma компании Google DeepMind (Gemma Team et al., 2024). CodeGemma — это семейство легких современных открытых моделей, созданных на основе тех же исследований и технологий, которые использовались при создании моделей Gemini.
Продолжая предварительно обученные модели Gemma, модели CodeGemma дополнительно обучаются на более чем 500–1000 миллиардах токенов, состоящих в основном из кода, с использованием той же архитектуры, что и семейство моделей Gemma. В результате модели CodeGemma обеспечивают современную производительность кода как при выполнении задач завершения, так и при генерации, сохраняя при этом глубокие навыки понимания и рассуждения в масштабе.
CodeGemma имеет 3 варианта:
- Предварительно обученная модель кода 7B
- Модель кода 7B, настроенная с помощью инструкций
- Модель 2B, специально обученная для заполнения кода и открытой генерации.
В этом руководстве рассказывается, как использовать модель CodeGemma с Flax для выполнения задачи завершения кода.
Настраивать
1. Настройте доступ Kaggle для CodeGemma.
Чтобы выполнить это руководство, сначала необходимо следовать инструкциям по установке на странице Gemma setup , в которых показано, как сделать следующее:
- Получите доступ к CodeGemma на kaggle.com .
- Выберите среду выполнения Colab с достаточными ресурсами ( у графического процессора T4 недостаточно памяти, вместо этого используйте TPU v2 ) для запуска модели CodeGemma.
- Создайте и настройте имя пользователя Kaggle и ключ API.
После завершения настройки Gemma перейдите к следующему разделу, где вы установите переменные среды для вашей среды Colab.
2. Установите переменные среды
Установите переменные среды для KAGGLE_USERNAME
и KAGGLE_KEY
. При появлении запроса «Предоставить доступ?» сообщения, согласитесь предоставить секретный доступ.
import os
from google.colab import userdata # `userdata` is a Colab API.
os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')
3. Установите библиотеку gemma
Бесплатного аппаратного ускорения Colab в настоящее время недостаточно для работы этого ноутбука. Если вы используете Colab Pay As You Go или Colab Pro , нажмите «Редактировать» > «Настройки ноутбука» > «Выбрать графический процессор A100» > «Сохранить» , чтобы включить аппаратное ускорение.
Далее вам необходимо установить библиотеку gemma
Google DeepMind с github.com/google-deepmind/gemma
. Если вы получаете сообщение об ошибке «преобразователь зависимостей pip», обычно вы можете игнорировать его.
pip install -q git+https://github.com/google-deepmind/gemma.git
4. Импортируйте библиотеки
В этом блокноте используется Gemma (которая использует Flax для построения слоев нейронной сети) и SentencePiece (для токенизации).
import os
from gemma import params as params_lib
from gemma import sampler as sampler_lib
from gemma import transformer as transformer_lib
import sentencepiece as spm
Загрузите модель CodeGemma
Загрузите модель CodeGemma с помощью kagglehub.model_download
, который принимает три аргумента:
-
handle
: ручка модели от Kaggle. -
path
: (Необязательная строка) Локальный путь -
force_download
: (Необязательное логическое значение) Принудительно повторно загрузить модель.
GEMMA_VARIANT = '2b-pt' # @param ['2b-pt', '7b-it', '7b-pt', '1.1-2b-pt', '1.1-7b-it'] {type:"string"}
import kagglehub
GEMMA_PATH = kagglehub.model_download(f'google/codegemma/flax/{GEMMA_VARIANT}')
Warning: Looks like you're using an outdated `kagglehub` version, please consider updating (latest version: 0.2.7) Downloading from https://www.kaggle.com/api/v1/models/google/codegemma/flax/2b-pt/3/download... 100%|██████████| 3.67G/3.67G [00:22<00:00, 173MB/s] Extracting model files...
print('GEMMA_PATH:', GEMMA_PATH)
GEMMA_PATH: /root/.cache/kagglehub/models/google/codegemma/flax/2b-pt/3
Проверьте расположение весов модели и токенизатора, затем установите переменные пути. Каталог токенизатора будет находиться в основном каталоге, в который вы загрузили модель, а веса модели будут находиться в подкаталоге. Например:
- Файл токенизатора
spm.model
будет находиться в/LOCAL/PATH/TO/codegemma/flax/2b-pt/3
- Контрольная точка модели будет находиться в
/LOCAL/PATH/TO/codegemma/flax/2b-pt/3/2b-pt
CKPT_PATH = os.path.join(GEMMA_PATH, GEMMA_VARIANT[-5:])
TOKENIZER_PATH = os.path.join(GEMMA_PATH, 'spm.model')
print('CKPT_PATH:', CKPT_PATH)
print('TOKENIZER_PATH:', TOKENIZER_PATH)
CKPT_PATH: /root/.cache/kagglehub/models/google/codegemma/flax/2b-pt/3/2b-pt TOKENIZER_PATH: /root/.cache/kagglehub/models/google/codegemma/flax/2b-pt/3/spm.model
Выполнить выборку/вывод
Загрузите и отформатируйте контрольную точку модели CodeGemma с помощью метода gemma.params.load_and_format_params
:
params = params_lib.load_and_format_params(CKPT_PATH)
Загрузите токенизатор CodeGemma, созданный с помощью sentencepiece.SentencePieceProcessor
:
vocab = spm.SentencePieceProcessor()
vocab.Load(TOKENIZER_PATH)
True
Чтобы автоматически загрузить правильную конфигурацию из контрольной точки модели CodeGemma, используйте gemma.transformer.TransformerConfig
. Аргумент cache_size
— это количество временных шагов в кеше CodeGemma Transformer
. После этого создайте экземпляр модели CodeGemma как model_2b
с помощью gemma.transformer.Transformer
(который наследуется от flax.linen.Module
).
transformer_config = transformer_lib.TransformerConfig.from_params(
params,
cache_size=1024
)
transformer = transformer_lib.Transformer(config=transformer_config)
Создайте sampler
с помощью gemma.sampler.Sampler
. Он использует контрольную точку модели CodeGemma и токенизатор.
sampler = sampler_lib.Sampler(
transformer=transformer,
vocab=vocab,
params=params['transformer']
)
Создайте несколько переменных для представления токенов заполнения посередине (fim) и создайте несколько вспомогательных функций для форматирования приглашения и сгенерированного вывода.
Например, давайте посмотрим на следующий код:
def function(string):
assert function('asdf') == 'fdsa'
Мы хотели бы заполнить function
так, чтобы утверждение выполнялось True
. В этом случае префикс будет таким:
"def function(string):\n"
И суффикс будет:
"assert function('asdf') == 'fdsa'"
Затем мы форматируем это в приглашение как PREFIX-SUFFIX-MIDDLE (средний раздел, который необходимо заполнить, всегда находится в конце приглашения):
"<|fim_prefix|>def function(string):\n<|fim_suffix|>assert function('asdf') == 'fdsa'<|fim_middle|>"
# In the context of a code editor,
# the cursor is the location where the text will be inserted
BEFORE_CURSOR = "<|fim_prefix|>"
AFTER_CURSOR = "<|fim_suffix|>"
AT_CURSOR = "<|fim_middle|>"
FILE_SEPARATOR = "<|file_separator|>"
def format_completion_prompt(before, after):
print(f"\nORIGINAL PROMPT:\n{before}{after}")
prompt = f"{BEFORE_CURSOR}{before}{AFTER_CURSOR}{after}{AT_CURSOR}"
print(f"\nFORMATTED PROMPT:\n{repr(prompt)}")
return prompt
def format_generated_output(before, after, output):
print(f"\nGENERATED OUTPUT:\n{repr(output)}")
formatted_output = f"{before}{output.replace(FILE_SEPARATOR, '')}{after}"
print(f"\nFILL-IN COMPLETION:\n{formatted_output}")
return formatted_output
Создайте подсказку и выполните вывод. Укажите префикс before
текстом и суффикс after
текста и сгенерируйте форматированное приглашение с помощью вспомогательной функции format_completion prompt
.
Вы можете настроить total_generation_steps
(количество шагов, выполняемых при генерации ответа — в этом примере используется 100
для экономии памяти хоста).
before = "def function(string):\n"
after = "assert function('asdf') == 'fdsa'"
prompt = format_completion_prompt(before, after)
output = sampler(
[prompt],
total_generation_steps=100,
).text
formatted_output = format_generated_output(before, after, output[0])
ORIGINAL PROMPT: def function(string): assert function('asdf') == 'fdsa' FORMATTED PROMPT: "<|fim_prefix|>def function(string):\n<|fim_suffix|>assert function('asdf') == 'fdsa'<|fim_middle|>" GENERATED OUTPUT: ' return string[::-1]\n\n<|file_separator|>' FILL-IN COMPLETION: def function(string): return string[::-1] assert function('asdf') == 'fdsa'
before = "import "
after = """if __name__ == "__main__":\n sys.exit(0)"""
prompt = format_completion_prompt(before, after)
output = sampler(
[prompt],
total_generation_steps=100,
).text
formatted_output = format_generated_output(before, after, output[0])
ORIGINAL PROMPT: import if __name__ == "__main__": sys.exit(0) FORMATTED PROMPT: '<|fim_prefix|>import <|fim_suffix|>if __name__ == "__main__":\n sys.exit(0)<|fim_middle|>' GENERATED OUTPUT: 'sys\n<|file_separator|>' FILL-IN COMPLETION: import sys if __name__ == "__main__": sys.exit(0)
before = """import numpy as np
def reflect(matrix):
# horizontally reflect a matrix
"""
after = ""
prompt = format_completion_prompt(before, after)
output = sampler(
[prompt],
total_generation_steps=100,
).text
formatted_output = format_generated_output(before, after, output[0])
ORIGINAL PROMPT: import numpy as np def reflect(matrix): # horizontally reflect a matrix FORMATTED PROMPT: '<|fim_prefix|>import numpy as np\ndef reflect(matrix):\n # horizontally reflect a matrix\n<|fim_suffix|><|fim_middle|>' GENERATED OUTPUT: ' return np.flip(matrix, axis=1)\n<|file_separator|>' FILL-IN COMPLETION: import numpy as np def reflect(matrix): # horizontally reflect a matrix return np.flip(matrix, axis=1)
Узнать больше
- Вы можете узнать больше о библиотеке
gemma
Google DeepMind на GitHub , которая содержит строки документации модулей, которые вы использовали в этом руководстве, таких какgemma.params
,gemma.transformer
иgemma.sampler
. - Следующие библиотеки имеют собственные сайты документации: core JAX , Flax и Orbax .
- Документацию по токенизатору/детокенизатору
sentencepiece
можно найти в репозитории Googlesentencepiece
на GitHub . - Документацию
kagglehub
можно найтиREADME.md
в репозиторииkagglehub
на GitHub . - Узнайте, как использовать модели Gemma с Google Cloud Vertex AI .
- Если вы используете Google Cloud TPU (v3-8 и новее), обязательно также обновите последнюю версию пакета
jax[tpu]
(!pip install -U jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
), перезапустите среду выполнения и проверьте соответствие версийjax
иjaxlib
(!pip list | grep jax
). Это может предотвратитьRuntimeError
, которая может возникнуть из-за несоответствия версийjaxlib
иjax
. Дополнительные инструкции по установке JAX см. в документации JAX .