Вывод с помощью Gemma с использованием JAX и Flax

Посмотреть на ai.google.dev Запустить в Google Colab Открыть в Vertex AI Посмотреть исходный код на GitHub

Обзор

Gemma — это семейство легких современных моделей открытого большого языка, основанных на исследованиях и технологиях Google DeepMind Gemini. В этом руководстве показано, как выполнить базовую выборку/вывод с помощью модели Gemma 2B Instruct с использованием библиотеки gemma Google DeepMind , написанной с использованием JAX (библиотека высокопроизводительных численных вычислений), Flax (библиотека нейронных сетей на основе JAX), Orbax (библиотека высокопроизводительных численных вычислений). библиотека на основе JAX для обучения таких утилит, как создание контрольных точек) и SentencePiece (библиотека токенизатора/детокенизатора). Хотя лен не используется напрямую в этой тетради, лен использовался для создания Джеммы.

Этот ноутбук может работать на Google Colab с бесплатным графическим процессором T4 (перейдите в «Правка» > «Настройки ноутбука» > в разделе «Аппаратный ускоритель» выберите «T4 GPU »).

Настраивать

1. Настройте доступ Kaggle для Джеммы.

Чтобы выполнить это руководство, сначала необходимо следовать инструкциям по установке на странице Gemma setup , в которых показано, как сделать следующее:

  • Получите доступ к Джемме на kaggle.com .
  • Выберите среду выполнения Colab с достаточными ресурсами для запуска модели Gemma.
  • Создайте и настройте имя пользователя 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. Чтобы включить аппаратное ускорение, нажмите «Редактировать» > «Настройки ноутбука» > «Выбрать графический процессор T4» > «Сохранить» .

Далее вам необходимо установить библиотеку gemma Google DeepMind с github.com/google-deepmind/gemma . Если вы получаете сообщение об ошибке «преобразователь зависимостей pip», обычно вы можете игнорировать его.

pip install -q git+https://github.com/google-deepmind/gemma.git

Загрузите и подготовьте модель Джеммы.

  1. Загрузите модель Gemma с помощью kagglehub.model_download , который принимает три аргумента:
  • handle : ручка модели от Kaggle.
  • path : (Необязательная строка) Локальный путь
  • force_download : (Необязательное логическое значение) Принудительно повторно загрузить модель.
GEMMA_VARIANT = '2b-it' # @param ['2b', '2b-it'] {type:"string"}

import kagglehub

GEMMA_PATH = kagglehub.model_download(f'google/gemma/flax/{GEMMA_VARIANT}')

Downloading from https://www.kaggle.com/api/v1/models/google/gemma/flax/2b-it/2/download...
100%|██████████| 3.67G/3.67G [00:35<00:00, 110MB/s]
Extracting model files...

print('GEMMA_PATH:', GEMMA_PATH)

GEMMA_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2
  1. Проверьте расположение весов модели и токенизатора, затем установите переменные пути. Каталог токенизатора будет находиться в основном каталоге, в который вы загрузили модель, а веса модели будут находиться в подкаталоге. Например:
  • Файл tokenizer.model будет находиться в /LOCAL/PATH/TO/gemma/flax/2b-it/2 ).
  • Контрольная точка модели будет находиться в /LOCAL/PATH/TO/gemma/flax/2b-it/2/2b-it ).
CKPT_PATH = os.path.join(GEMMA_PATH, GEMMA_VARIANT)
TOKENIZER_PATH = os.path.join(GEMMA_PATH, 'tokenizer.model')
print('CKPT_PATH:', CKPT_PATH)
print('TOKENIZER_PATH:', TOKENIZER_PATH)

CKPT_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/2b-it
TOKENIZER_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/tokenizer.model

Выполнить выборку/вывод

  1. Загрузите и отформатируйте контрольную точку модели Gemma с помощью метода gemma.params.load_and_format_params :
from gemma import params as params_lib

params = params_lib.load_and_format_params(CKPT_PATH)
  1. Загрузите токенизатор Gemma, созданный с помощью sentencepiece.SentencePieceProcessor :
import sentencepiece as spm

vocab = spm.SentencePieceProcessor()
vocab.Load(TOKENIZER_PATH)

True
  1. Чтобы автоматически загрузить правильную конфигурацию из контрольной точки модели Gemma, используйте gemma.transformer.TransformerConfig . Аргумент cache_size — это количество временных шагов в кеше Gemma Transformer . После этого создайте экземпляр модели Gemma как transformer с помощью gemma.transformer.Transformer (который наследуется от flax.linen.Module ).
from gemma import transformer as transformer_lib

transformer_config = transformer_lib.TransformerConfig.from_params(
    params=params,
    cache_size=1024
)

transformer = transformer_lib.Transformer(transformer_config)
  1. Создайте sampler с помощью gemma.sampler.Sampler поверх контрольной точки/весов модели Gemma и токенизатора:
from gemma import sampler as sampler_lib

sampler = sampler_lib.Sampler(
    transformer=transformer,
    vocab=vocab,
    params=params['transformer'],
)
  1. Напишите запрос в input_batch и выполните вывод. Вы можете настроить total_generation_steps (количество шагов, выполняемых при генерации ответа — в этом примере используется 100 для экономии памяти хоста).
prompt = [
    "\n# What is the meaning of life?",
]

reply = sampler(input_strings=prompt,
                total_generation_steps=100,
                )

for input_string, out_string in zip(prompt, reply.text):
    print(f"Prompt:\n{input_string}\nOutput:\n{out_string}")

Prompt:

# What is the meaning of life?
Output:


The question of what the meaning of life is one that has occupied the minds of philosophers, theologians, and individuals for centuries. There is no single, universally accepted answer, but there are many different perspectives on this complex and multifaceted question.

**Some common perspectives on the meaning of life include:**

* **Biological perspective:** From a biological standpoint, the meaning of life is to survive and reproduce.
* **Existential perspective:** Existentialists believe that life is not inherently meaningful and that
  1. (Необязательно) Запустите эту ячейку, чтобы освободить память, если вы завершили записную книжку и хотите попробовать еще одну подсказку. После этого вы можете снова создать экземпляр sampler на шаге 3, а затем настроить и запустить приглашение на шаге 4.
del sampler

Узнать больше