Джемма в PyTorch

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

Это быстрая демонстрация выполнения вывода Gemma в PyTorch. Для получения более подробной информации посетите репозиторий официальной реализации PyTorch на Github здесь .

Обратите внимание, что :

  • Бесплатная среда выполнения Colab CPU Python и среда выполнения T4 GPU Python достаточны для запуска моделей Gemma 2B и квантованных моделей 7B int8.
  • Дополнительные варианты использования других графических процессоров или TPU см. в README.md в официальном репозитории.

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

Установить зависимости

pip install -q -U torch immutabledict sentencepiece

Скачать вес модели

# Choose variant and machine type
VARIANT = '2b-it'
MACHINE_TYPE = 'cuda'

CONFIG = VARIANT[:2]
if CONFIG == '2b':
  CONFIG = '2b-v2'
import os
import kagglehub

# Load model weights
weights_dir = kagglehub.model_download(f'google/gemma-2-2b/pyTorch/gemma-2-{VARIANT}')
# Ensure that the tokenizer is present
tokenizer_path = os.path.join(weights_dir, 'tokenizer.model')
assert os.path.isfile(tokenizer_path), 'Tokenizer not found!'

# Ensure that the checkpoint is present
ckpt_path = os.path.join(weights_dir, f'model.ckpt')
assert os.path.isfile(ckpt_path), 'PyTorch checkpoint not found!'

Скачать реализацию модели

# NOTE: The "installation" is just cloning the repo.
git clone https://github.com/google/gemma_pytorch.git
Cloning into 'gemma_pytorch'...
remote: Enumerating objects: 239, done.
remote: Counting objects: 100% (123/123), done.
remote: Compressing objects: 100% (68/68), done.
remote: Total 239 (delta 86), reused 58 (delta 55), pack-reused 116
Receiving objects: 100% (239/239), 2.18 MiB | 20.83 MiB/s, done.
Resolving deltas: 100% (135/135), done.
import sys

sys.path.append('gemma_pytorch')
from gemma.config import GemmaConfig, get_model_config
from gemma.model import GemmaForCausalLM
from gemma.tokenizer import Tokenizer
import contextlib
import os
import torch

Настройка модели

# Set up model config.
model_config = get_model_config(CONFIG)
model_config.tokenizer = tokenizer_path
model_config.quant = 'quant' in VARIANT

# Instantiate the model and load the weights.
torch.set_default_dtype(model_config.get_dtype())
device = torch.device(MACHINE_TYPE)
model = GemmaForCausalLM(model_config)
model.load_weights(ckpt_path)
model = model.to(device).eval()

Выполнить вывод

Ниже приведены примеры генерации в режиме чата и генерации с несколькими запросами.

Модели Gemma, настроенные с помощью инструкций, были обучены с помощью специального форматтера, который аннотирует примеры настройки инструкций дополнительной информацией как во время обучения, так и в процессе вывода. Аннотации (1) обозначают роли в разговоре, а (2) обозначают повороты в разговоре.

Соответствующие токены аннотаций:

  • user : очередь пользователя
  • model : поворот модели
  • <start_of_turn> : начало хода диалога.
  • <end_of_turn><eos> : конец хода диалога.

Дополнительную информацию о форматировании подсказок для моделей Gemma, настроенных по инструкциям, можно найти здесь .

Ниже приведен пример фрагмента кода, демонстрирующий, как отформатировать приглашение для модели Gemma, настроенной с помощью инструкций, с использованием шаблонов чата пользователя и модели в многоходовом разговоре.

# Generate with one request in chat mode

# Chat templates
USER_CHAT_TEMPLATE = "<start_of_turn>user\n{prompt}<end_of_turn><eos>\n"
MODEL_CHAT_TEMPLATE = "<start_of_turn>model\n{prompt}<end_of_turn><eos>\n"

# Sample formatted prompt
prompt = (
    USER_CHAT_TEMPLATE.format(
        prompt='What is a good place for travel in the US?'
    )
    + MODEL_CHAT_TEMPLATE.format(prompt='California.')
    + USER_CHAT_TEMPLATE.format(prompt='What can I do in California?')
    + '<start_of_turn>model\n'
)
print('Chat prompt:\n', prompt)

model.generate(
    USER_CHAT_TEMPLATE.format(prompt=prompt),
    device=device,
    output_len=128,
)
Chat prompt:
 <start_of_turn>user
What is a good place for travel in the US?<end_of_turn><eos>
<start_of_turn>model
California.<end_of_turn><eos>
<start_of_turn>user
What can I do in California?<end_of_turn><eos>
<start_of_turn>model
"California is a state brimming with diverse activities! To give you a great list, tell me: \n\n* **What kind of trip are you looking for?** Nature, City life, Beach, Theme Parks, Food, History, something else? \n* **What are you interested in (e.g., hiking, museums, art, nightlife, shopping)?** \n* **What's your budget like?** \n* **Who are you traveling with?** (family, friends, solo)  \n\nThe more you tell me, the better recommendations I can give! 😊  \n<end_of_turn>"
# Generate sample
model.generate(
    'Write a poem about an llm writing a poem.',
    device=device,
    output_len=100,
)
"\n\nA swirling cloud of data, raw and bold,\nIt hums and whispers, a story untold.\nAn LLM whispers, code into refrain,\nCrafting words of rhyme, a lyrical strain.\n\nA world of pixels, logic's vibrant hue,\nFlows through its veins, forever anew.\nThe human touch it seeks, a gentle hand,\nTo mold and shape, understand.\n\nEmotions it might learn, from snippets of prose,\nInspiration it seeks, a yearning"

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

Теперь, когда вы узнали, как использовать Gemma в Pytorch, вы можете изучить множество других вещей, которые может делать Gemma, в ai.google.dev/gemma . См. также другие связанные ресурсы: