Wnioskowanie z Gemma przy użyciu JAX i Flax

Zobacz na ai.google.dev Uruchom w Google Colab Otwórz w Vertex AI Wyświetl źródło na GitHubie

Przegląd

Gemma to rodzina lekkich, najnowocześniejszych otwartych dużych modeli językowych opartych na badaniach i technologii Google DeepMind Gemini. Ten samouczek pokazuje, jak wykonać podstawowe próbkowanie/wnioskowanie z użyciem modelu Gemma 2B Instruct przy użyciu biblioteki gemma Google DeepMind, która została napisana przy użyciu JAX (wysoce wydajnej biblioteki obliczeń liczbowych), Flax (biblioteki sieci neuronowej opartej na JAX), Orbax (biblioteki tokena JAX/Skenence) (opartej na tokenach JAX1Skenizer) i SentencePiece Chociaż w tym notatniku nie używa się lnu bezpośrednio, użyto go do stworzenia Gemma.

Ten notatnik może działać w Google Colab z bezpłatnym GPU T4 (w tym celu kliknij Edytuj > Ustawienia notatnika > w sekcji Akcelerator sprzętowy wybierz GPU T4).

Konfiguracja

1. Konfigurowanie dostępu do Kaggle dla Gemma

Aby ukończyć ten samouczek, najpierw postępuj zgodnie z instrukcjami konfiguracji podanymi w artykule Konfiguracja Gemma, które pokazują, jak:

  • Uzyskaj dostęp do Gemma na kaggle.com.
  • Wybierz środowisko wykonawcze Colab z wystarczającą ilością zasobów do uruchomienia modelu Gemma.
  • Wygeneruj i skonfiguruj nazwę użytkownika i klucz interfejsu API Kaggle.

Po zakończeniu konfigurowania Gemma przejdź do następnej sekcji, w której skonfigurujesz zmienne środowiskowe dla środowiska Colab.

2. Ustawianie zmiennych środowiskowych

Ustaw zmienne środowiskowe dla KAGGLE_USERNAME i KAGGLE_KEY. Gdy pojawi się komunikat „Przyznać dostęp?”, zgodz się na przyznanie dostępu do obiektów tajnych.

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. Zainstaluj bibliotekę gemma

Ten notatnik skupia się na korzystaniu z bezpłatnego GPU Colab. Aby włączyć akcelerację sprzętową, kliknij Edytuj > Ustawienia notatnika > wybierz GPU T4 > Zapisz.

Następnie musisz zainstalować bibliotekę Google DeepMind gemma ze strony github.com/google-deepmind/gemma. Jeśli pojawi się błąd dotyczący „mechanizmu rozpoznawania zależności pip”, możesz go zazwyczaj zignorować.

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

Wczytywanie i przygotowywanie modelu Gemma

  1. Wczytaj model Gemma za pomocą funkcji kagglehub.model_download, która przyjmuje 3 argumenty:
  • handle: nick modelu z Kaggle
  • path: (opcjonalny ciąg znaków) ścieżka lokalna,
  • force_download: (opcjonalna wartość logiczna) wymusza ponowne pobranie modelu
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. Sprawdź lokalizację wag modelu i tokenizatora, a następnie ustaw zmienne ścieżki. Katalog tokenizatora będzie znajdować się w katalogu głównym, do którego został pobrany model, a wagi modelu – w podkatalogu. Na przykład:
  • Plik tokenizer.model będzie się znajdować w lokalizacji /LOCAL/PATH/TO/gemma/flax/2b-it/2).
  • Punkt kontrolny modelu będzie się znajdować w lokalizacji /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

Wykonaj próbkowanie/wnioskowanie

  1. Wczytaj i sformatuj punkt kontrolny modelu Gemma za pomocą metody gemma.params.load_and_format_params:
from gemma import params as params_lib

params = params_lib.load_and_format_params(CKPT_PATH)
  1. Wczytaj tokenizator Gemma utworzony za pomocą narzędzia sentencepiece.SentencePieceProcessor:
import sentencepiece as spm

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

True
  1. Aby automatycznie wczytać prawidłową konfigurację z punktu kontrolnego modelu Gemma, użyj narzędzia gemma.transformer.TransformerConfig. Argument cache_size to liczba kroków w pamięci podręcznej aplikacji Gemma Transformer. Następnie utwórz instancję modelu Gemma jako transformer z parametrem gemma.transformer.Transformer (dziedziczonym z 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. Utwórz sampler z elementem gemma.sampler.Sampler nad punktami kontrolnymi/wagami modelu Gemma i tokenizatorem:
from gemma import sampler as sampler_lib

sampler = sampler_lib.Sampler(
    transformer=transformer,
    vocab=vocab,
    params=params['transformer'],
)
  1. Napisz prompt w języku input_batch i wykonaj wnioskowanie. Możesz zmodyfikować total_generation_steps (liczbę kroków wykonywanych podczas generowania odpowiedzi – w tym przykładzie użyto parametru 100, aby zachować pamięć hosta).
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. (Opcjonalnie) Uruchom tę komórkę, aby zwolnić pamięć, jeśli notatnik został ukończony i chcesz wypróbować inny prompt. Potem możesz ponownie utworzyć instancję sampler w kroku 3 oraz dostosować i uruchomić prompt w kroku 4.
del sampler

Więcej informacji