 ai.google.dev'de görüntüleyin
|
 Google Colab'de çalıştır
|
 Vertex AI'da aç
|
 Kaynağı GitHub'da görüntüleyin
|
Genel bakış
Gemma, Google DeepMind Gemini araştırmasına ve teknolojisine dayanan, hafif ve son teknoloji ürünü açık büyük dil modellerinden oluşan bir ailedir. Bu eğitim; Google DeepMind'ın gemma kitaplığı, JAX (yüksek performanslı sayısal bilgi işlem kitaplığı), Flax (JAX tabanlı nöral ağ kitaplığı), Chex (güvenilir JAX tabanlı JAX kodu kitaplığı) kullanarak İngilizce-Fransızca çeviri görevi için Gemma 2B Talimatlar modelinin ince ayarını nasıl yapacağınızı gösterir. Bu not defterinde doğrudan kullanılmasa da, Gemma'yı oluşturmak için Flax kullanıldı.
gemma kitaplığı JAX, Flax, Orbax (kontrol noktası gibi eğitim yardımcı programlarına yönelik JAX tabanlı bir kitaplık) ve SentencePiece (bir jeton belirleyici/detokenizer kitaplığı) ile yazılmıştır.
Kurulum
1. Gemma için Kaggle erişimini ayarlama
Bu eğiticiyi tamamlamak için önce Gemma kurulumu sayfasındaki kurulum talimatlarını uygulamanız gerekir. Bu talimatlarda, aşağıdakilerin nasıl yapılacağı gösterilir:
- kaggle.com adresinden Gemma'ya erişin.
- Gemma modelini çalıştırmak için yeterli kaynağa sahip bir Colab çalışma zamanı seçin.
- Bir Kaggle kullanıcı adı ve API anahtarı oluşturun ve yapılandırın.
Gemma kurulumunu tamamladıktan sonra, Colab ortamınız için ortam değişkenlerini belirleyeceğiniz bir sonraki bölüme geçin.
2. Ortam değişkenlerini ayarlama
KAGGLE_USERNAME ve KAGGLE_KEY için ortam değişkenlerini ayarlayın. "Erişim izni verilsin mi?" iletileriyle karşılaştığınızda gizli erişim sağlamayı kabul edin.
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 kitaplığını yükle
Ücretsiz Colab donanım hızlandırma özelliği şu anda bu not defterini çalıştırmak için yetersiz. Colab Pay As You Go veya Colab Pro kullanıyorsanız Düzenle > Not defteri ayarları'nı tıklayın > A100 GPU > Kaydet'i seçerek donanım hızlandırmayı etkinleştirin.
Ardından, github.com/google-deepmind/gemma kaynağından Google DeepMind gemma kitaplığını yüklemeniz gerekir. "pip'in bağımlılık çözümleyicisi" ile ilgili bir hata alırsanız genellikle bu hatayı yoksayabilirsiniz.
pip install -q git+https://github.com/google-deepmind/gemma.git
4. Kitaplıkları içe aktar
Bu not defterinde Flax (nöral ağlar için), çekirdek JAX, SentencePiece (jetonleştirme için), Chex (güvenilir JAX kodu yazmaya yönelik yardımcı program kitaplığı) ve TensorFlow Veri Kümeleri kullanılır.
import os
import enum
import re
import string
import chex
import jax
import jax.numpy as jnp
import optax
import tensorflow as tf
import tensorflow_datasets as tfds
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
Gemma modelini yükleme
Üç bağımsız değişken alan kagglehub.model_download ile Gemma modelini yükleyin:
handle: Kaggle'daki model herkese açık kullanıcı adıpath: (İsteğe bağlı dize) Yerel yolforce_download: (İsteğe bağlı boole) Modeli yeniden indirmeye zorlar
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:26<00:00, 147MB/s] Extracting model files...
print('GEMMA_PATH:', GEMMA_PATH)
GEMMA_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2
Model ağırlıklarının ve belirteçleştiricinin konumunu kontrol edin, ardından yol değişkenlerini ayarlayın. Jeton oluşturucu dizini, modeli indirdiğiniz ana dizindeyken, model ağırlıkları bir alt dizinde yer alır. Örneğin:
tokenizer.modeldosyası/LOCAL/PATH/TO/gemma/flax/2b-it/2içinde olacaktır).- Model kontrol noktası
/LOCAL/PATH/TO/gemma/flax/2b-it/2/2b-itiçinde olacaktır).
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
MTNT veri kümesini ve Gemma belirteç oluşturucuyu yükleme ve hazırlama
TensorFlow Datasets'ten erişilebilen MTNT (Machine Translation of Noisy Text) veri kümesini kullanacaksınız.
MTNT veri kümesinin İngilizce-Fransızca veri kümesi bölümünü indirin ve ardından iki örneği inceleyin. Veri kümesindeki her örnek iki giriş içerir: src: orijinal İngilizce cümle ve dst: karşılık gelen Fransızca çeviri.
ds = tfds.load("mtnt/en-fr", split="train")
ds = ds.take(2)
ds = ds.as_numpy_iterator()
for idx, example in enumerate(ds):
print(f'Example {idx}:')
for key, val in example.items():
print(f'{key}: {val}')
print()
Downloading and preparing dataset 35.08 MiB (download: 35.08 MiB, generated: 11.33 MiB, total: 46.41 MiB) to /root/tensorflow_datasets/mtnt/en-fr/1.0.0... Dl Completed...: 0 url [00:00, ? url/s] Dl Size...: 0 MiB [00:00, ? MiB/s] Extraction completed...: 0 file [00:00, ? file/s] Generating splits...: 0%| | 0/3 [00:00<?, ? splits/s] Generating train examples...: 0%| | 0/35692 [00:00<?, ? examples/s] Shuffling /root/tensorflow_datasets/mtnt/en-fr/1.0.0.incomplete6YJMND/mtnt-train.tfrecord*...: 0%| … Generating test examples...: 0%| | 0/1020 [00:00<?, ? examples/s] Shuffling /root/tensorflow_datasets/mtnt/en-fr/1.0.0.incomplete6YJMND/mtnt-test.tfrecord*...: 0%| |… Generating valid examples...: 0%| | 0/811 [00:00<?, ? examples/s] Shuffling /root/tensorflow_datasets/mtnt/en-fr/1.0.0.incomplete6YJMND/mtnt-valid.tfrecord*...: 0%| … Dataset mtnt downloaded and prepared to /root/tensorflow_datasets/mtnt/en-fr/1.0.0. Subsequent calls will reuse this data. Example 0: dst: b'Le groupe de " toutes les \xc3\xa9toiles potentielles de la conf\xc3\xa9rence de l\'Est mais qui ne s\'en sortent pas dans le groupe de l\'Ouest ".' src: b'The group of \xe2\x80\x9ceastern conference potential all stars but not making it in the West\xe2\x80\x9d group.' Example 1: dst: b"Kameron est-elle un peu aigrie de son manque de temps \xc3\xa0 l'\xc3\xa9cran ?" src: b'Is Kameron a Little Salty About Her Lack of Air Time?'
sentencepiece.SentencePieceProcessor kullanılarak oluşturulan Gemma belirteç oluşturucuyu yükleyin:
vocab = spm.SentencePieceProcessor()
vocab.Load(TOKENIZER_PATH)
True
İngilizce-Fransızca çeviri görevi için SentencePieceProcessor öğesini özelleştirin. Gemma modelinin İngilizce kısmına ince ayar yapacağınız için birkaç ayarlama yapmanız gerekir. Örneğin:
Giriş öneki: Her girişe ortak bir ön ek eklemek, çeviri görevine işaret eder. Örneğin,
Translate this into French: [INPUT_SENTENCE]gibi ön eke sahip bir istem kullanabilirsiniz.Çeviri başlangıç son eki: Her istemin sonuna bir sonek eklenmesi, Gemma modeline çeviri işleminin tam olarak ne zaman başlayacağını bildirir. Yeni bir satır yeterli olacaktır.
Dil modeli jetonları: Gemma modelleri, her dizinin başında bir "dizinin başlangıcı" jetonu bulunmasını bekler. Bu nedenle, her eğitim örneğinin sonuna bir "dizi sonu" jetonu eklemek yeterli olacaktır.
SentencePieceProcessoröğesinin çevresinde aşağıdaki gibi bir özel sarmalayıcı oluşturun:
class GemmaTokenizer:
def __init__(self,
spm_processor: spm.SentencePieceProcessor):
self._spm_processor = spm_processor
@property
def pad_id(self) -> int:
"""Fast access to the pad ID."""
return self._spm_processor.pad_id()
def tokenize(self,
example: str | bytes,
prefix: str = '',
suffix: str = '',
add_eos: bool = True) -> jax.Array:
"""
The tokenization function.
Args:
example: Input string to tokenize.
prefix: Prefix to add to the input string.
suffix: Suffix to add to the input string.
add_eos: If True, add an "end of sentence" token at the end of the output
sequence.
Returns:
Tokens corresponding to the input string.
"""
int_list = [self._spm_processor.bos_id()]
int_list.extend(self._spm_processor.EncodeAsIds(prefix + example + suffix))
if add_eos:
int_list.append(self._spm_processor.eos_id())
return jnp.array(int_list, dtype=jnp.int32)
def tokenize_tf_op(self,
str_tensor: tf.Tensor,
prefix: str = '',
suffix: str = '',
add_eos: bool = True) -> tf.Tensor:
"""A TensorFlow operator for the tokenize function."""
encoded = tf.numpy_function(
self.tokenize,
[str_tensor, prefix, suffix, add_eos],
tf.int32)
encoded.set_shape([None])
return encoded
def to_string(self, tokens: jax.Array) -> str:
"""Convert an array of tokens to a string."""
return self._spm_processor.EncodeIds(tokens.tolist())
Yeni özel GemmaTokenizer öğenizi örnekleyerek ve ardından bu örneği MTNT veri kümesinin küçük bir örneğine uygulayarak deneyin:
tokenizer = GemmaTokenizer(vocab)
def tokenize_source(tokenizer, example: tf.Tensor):
return tokenizer.tokenize_tf_op(example,
prefix='Translate this into French:\n',
suffix='\n',
add_eos=False)
def tokenize_destination(tokenizer, example: tf.Tensor):
return tokenizer.tokenize_tf_op(example,
add_eos=True)
ds = tfds.load("mtnt/en-fr",split="train")
ds = ds.take(2)
ds = ds.map(lambda x: {'src': tokenize_source(tokenizer, x['src']),
'dst': tokenize_destination(tokenizer, x['dst'])})
ds = ds.as_numpy_iterator()
for idx, example in enumerate(ds):
print(f'Example {idx}:')
for key, val in example.items():
print(f'{key}: {val}')
print()
Example 0:
src: [ 2 49688 736 1280 6987 235292 108 651 2778 576
1080 104745 11982 5736 832 8995 901 780 3547 665
575 573 4589 235369 2778 235265 108]
dst: [ 2 2025 29653 581 664 16298 1437 55563 41435 7840
581 683 111452 581 533 235303 9776 4108 2459 679
485 235303 479 6728 579 1806 2499 709 29653 581
533 235303 101323 16054 1]
Example 1:
src: [ 2 49688 736 1280 6987 235292 108 2437 87150 477
476 11709 230461 8045 3636 40268 576 4252 4897 235336
108]
dst: [ 2 213606 477 1455 235290 3510 748 8268 191017 2809
581 2032 69972 581 11495 1305 533 235303 65978 1654
1]
Tüm MTNT veri kümesi için bir veri yükleyici oluşturun:
@chex.dataclass(frozen=True)
class TrainingInput:
# Input tokens provided to the model.
input_tokens: jax.Array
# A mask that determines which tokens contribute to the target loss
# calculation.
target_mask: jax.Array
class DatasetSplit(enum.Enum):
TRAIN = 'train'
VALIDATION = 'valid'
class MTNTDatasetBuilder:
"""The dataset builder for the MTNT dataset."""
N_ITEMS = {DatasetSplit.TRAIN: 35_692,
DatasetSplit.VALIDATION: 811}
BUFFER_SIZE_SHUFFLE = 10_000
TRANSLATION_PREFIX = 'Translate this into French:\n'
TRANSLATION_SUFFIX = '\n'
def __init__(self,
tokenizer : GemmaTokenizer,
max_seq_len: int):
"""Constructor.
Args:
tokenizer: Gemma tokenizer to use.
max_seq_len: size of each sequence in a given batch.
"""
self._tokenizer = tokenizer
self._base_data = {
DatasetSplit.TRAIN: tfds.load("mtnt/en-fr",split="train"),
DatasetSplit.VALIDATION: tfds.load("mtnt/en-fr",split="valid"),
}
self._max_seq_len = max_seq_len
def _tokenize_source(self, example: tf.Tensor):
"""Tokenization function for the source."""
return self._tokenizer.tokenize_tf_op(example,
prefix=self.TRANSLATION_PREFIX,
suffix=self.TRANSLATION_SUFFIX,
add_eos=False)
def _tokenize_destination(self, example: tf.Tensor):
"""Tokenization function for the French translation."""
return self._tokenizer.tokenize_tf_op(example,
add_eos=True)
def _pad_up_to_max_len(self,
input_tensor: tf.Tensor,
pad_value: int | bool,
) -> tf.Tensor:
"""Pad the given tensor up to sequence length of a batch."""
seq_len = tf.shape(input_tensor)[0]
to_pad = tf.maximum(self._max_seq_len - seq_len, 0)
return tf.pad(input_tensor,
[[0, to_pad]],
mode='CONSTANT',
constant_values=pad_value,
)
def _to_training_input(self,
src_tokens: jax.Array,
dst_tokens: jax.Array,
) -> TrainingInput:
"""Build a training input from a tuple of source and destination tokens."""
# The input sequence fed to the model is simply the concatenation of the
# source and the destination.
tokens = tf.concat([src_tokens, dst_tokens], axis=0)
# To prevent the model from updating based on the source (input)
# tokens, add a target mask to each input.
q_mask = tf.zeros_like(src_tokens, dtype=tf.bool)
a_mask = tf.ones_like(dst_tokens, dtype=tf.bool)
mask = tf.concat([q_mask, a_mask], axis=0)
# If the output tokens sequence is smaller than the target sequence size,
# then pad it with pad tokens.
tokens = self._pad_up_to_max_len(tokens, self._tokenizer.pad_id)
# Don't want to perform the backward pass on the pad tokens.
mask = self._pad_up_to_max_len(mask, False)
return TrainingInput(input_tokens=tokens, target_mask=mask)
def get_train_dataset(self, batch_size: int, num_epochs: int):
"""Build the training dataset."""
# Tokenize each sample.
ds = self._base_data[DatasetSplit.TRAIN].map(lambda x : (self._tokenize_source(x['src']),
self._tokenize_destination(x['dst'])))
# Convert the samples to training inputs.
ds = ds.map(lambda x, y: self._to_training_input(x, y))
# Remove the samples that are too long.
ds = ds.filter(lambda x: tf.shape(x.input_tokens)[0] <= self._max_seq_len)
# Shuffle the dataset.
ds = ds.shuffle(buffer_size=self.BUFFER_SIZE_SHUFFLE)
# Repeat if necessary.
ds = ds.repeat(num_epochs)
# Build batches.
ds = ds.batch(batch_size, drop_remainder=True)
return ds
def get_validation_dataset(self, batch_size: int):
"""Build the validation dataset."""
# Same steps as in `get_train_dataset`, but without shuffling and no repetition.
ds = self._base_data[DatasetSplit.VALIDATION].map(lambda x : (self._tokenize_source(x['src']),
self._tokenize_destination(x['dst'])))
ds = ds.map(lambda x, y: self._to_training_input(x, y))
ds = ds.filter(lambda x: tf.shape(x.input_tokens)[0] <= self._max_seq_len)
ds = ds.batch(batch_size, drop_remainder=True)
return ds
MTNTDatasetBuilder özelliğini denemek için özel GemmaTokenizer öğesini yeniden örnekleyin, ardından MTNT veri kümesine uygulayın ve iki örnek örnekleyin:
tokenizer = GemmaTokenizer(vocab)
dataset_builder = MTNTDatasetBuilder(tokenizer, max_seq_len=20)
ds = dataset_builder.get_train_dataset(3, 1)
ds = ds.take(2)
ds = ds.as_numpy_iterator()
for idx, example in enumerate(ds):
print(f'Example {idx}:')
for key, val in example.items():
print(f'{key}: {val}')
print()
WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class '__main__.TrainingInput'>
WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class '__main__.TrainingInput'>
WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class '__main__.TrainingInput'>
Example 0:
input_tokens: [[ 2 49688 736 1280 6987 235292 108 10924 665 12302
235341 108 2 4397 63011 1437 38696 1241 1 0]
[ 2 49688 736 1280 6987 235292 108 13835 1517 235265
108 2 69875 540 19713 235265 1 0 0 0]
[ 2 49688 736 1280 6987 235292 108 6956 1586 235297
235265 108 2 78368 1586 235297 235265 1 0 0]]
target_mask: [[False False False False False False False False False False False False
True True True True True True True False]
[False False False False False False False False False False False True
True True True True True False False False]
[False False False False False False False False False False False False
True True True True True True False False]]
Example 1:
input_tokens: [[ 2 49688 736 1280 6987 235292 108 18874 235341 108
2 115905 6425 1241 1 0 0 0 0 0]
[ 2 49688 736 1280 6987 235292 108 7574 3356 235341
108 2 7997 20707 1241 1 0 0 0 0]
[ 2 49688 736 1280 6987 235292 108 8703 665 235265
108 2 235338 235303 90006 20133 235265 1 0 0]]
target_mask: [[False False False False False False False False False False True True
True True True False False False False False]
[False False False False False False False False False False False True
True True True True False False False False]
[False False False False False False False False False False False True
True True True True True True False False]]
Modeli yapılandırma
Gemma modeline ince ayar yapmaya başlamadan önce modeli yapılandırmanız gerekir.
Öncelikle Gemma modeli kontrol noktasını gemma.params.load_and_format_params yöntemiyle yükleyin ve biçimlendirin:
params = params_lib.load_and_format_params(CKPT_PATH)
Gemma modeli kontrol noktasından doğru yapılandırmayı otomatik olarak yüklemek için gemma.transformer.TransformerConfig değerini kullanın. cache_size bağımsız değişkeni, Gemma Transformer önbelleğindeki adım sayısıdır. Daha sonra, Gemma modelini gemma.transformer.Transformer ile model_2b olarak örnekleyin (flax.linen.Module öğesinden devralır).
config_2b = transformer_lib.TransformerConfig.from_params(
params,
cache_size=30
)
model_2b = transformer_lib.Transformer(config=config_2b)
Modelde ince ayar yapma
Bu bölümde şunları yapacaksınız:
- İleriye doğru hesaplama ve kaybetme işlevini oluşturmak için
gemma.transformer.Transformersınıfını kullanın. - Jetonlar için konum ve dikkat maskesi vektörlerini oluşturma
- Flax ile bir eğitim adım işlevi oluşturun.
- Doğrulama adımını geriye dönük geçiş olmadan oluşturun.
- Eğitim döngüsünü oluşturun.
- Gemma modelinde ince ayar yapın.
gemma.transformer.Transformer sınıfını kullanarak ileriye doğru hesaplama ve kayıp işlevini tanımlayın. Gemma Transformer, flax.linen.Module'ten devralmaktadır ve iki temel yöntem sunar:
init: Modelin parametrelerini başlatır.apply: Belirli bir parametre grubunu kullanarak modelin__call__işlevini yürütür.Önceden eğitilmiş Gemma ağırlıklarıyla çalıştığınız için
initişlevini kullanmanıza gerek yoktur.
def forward_and_loss_fn(params,
*,
model: transformer_lib.Transformer,
input_tokens: jax.Array, # Shape [B, L]
input_mask: jax.Array, # Shape [B, L]
positions: jax.Array, # Shape [B, L]
attention_mask: jax.Array, # [B, L, L]
) -> jax.Array:
"""The forward pass and the loss function.
Args:
params: Model's input parameters.
model: The Gemma transformer model to call.
input_tokens: Input tokens sequence, shape [B, L].
input_mask: Tokens to ignore when computing the loss, shape [B, L].
positions: Relative position of each token, shape [B, L].
attention_mask: Input attention mask, shape [B, L].
Returns:
The softmax cross-entropy loss for the next-token prediction task.
"""
# The forward pass on the input data.
# No attention cache is needed here.
logits, _ = model.apply(
params,
input_tokens,
positions,
None, # Attention cache is None.
attention_mask,
)
# Exclude the last step as it does not appear in the targets.
logits = logits[0, :-1]
# Similarly, the first token cannot be predicted.
target_tokens = input_tokens[0, 1:]
target_mask = input_mask[0, 1:]
# Convert the target labels to one-hot encoded vectors.
one_hot = jax.nn.one_hot(target_tokens, logits.shape[-1])
# Don't update on unwanted tokens.
one_hot = one_hot * target_mask.astype(one_hot.dtype)[...,None]
# Define the normalization factor.
norm_factor = 1 / (jnp.sum(target_mask) + 1e-8)
# Return the negative log likelihood (NLL) loss.
return -jnp.sum(jax.nn.log_softmax(logits) * one_hot) * norm_factor
gemma.transformer.Transformer sınıfı için her girişle birlikte bir attention_mask ve positions vektörü gerekir. Bunları, Transformer.build_positions_from_mask ve Transformer.make_causal_attn_mask kullanan özel bir işlev oluşturarak oluşturabilirsiniz:
def get_attention_mask_and_positions(example: jax.Array,
pad_id : int,
)-> tuple[jax.Array, jax.Array]:
"""Builds the position and attention mask vectors from the given tokens."""
pad_mask = example != pad_id
current_token_position = transformer_lib.build_positions_from_mask(pad_mask)
attention_mask = transformer_lib.make_causal_attn_mask(pad_mask)
return current_token_position, attention_mask
Geriye doğru geçişi gerçekleştiren ve modelin parametrelerini uygun şekilde güncelleyen train_step işlevini oluşturun. Bunun için aşağıdaki örnekler verilebilir:
jax.value_and_grad, ileri ve geri geçişleri sırasında kayıp işlevini ve renk geçişlerini değerlendirmek içindir.optax.apply_updates, parametreleri güncellemek içindir.
def train_step(model: transformer_lib.Transformer,
params,
optimizer: optax.GradientTransformation,
opt_state: optax.OptState,
pad_id: int,
example: TrainingInput):
"""Train step.
Args:
model: The Gemma transformer model.
params: The model's input parameters.
optimizer: The Optax optimizer to use.
opt_state: The input optimizer's state.
pad_id: ID of the pad token.
example: Input batch.
Returns:
The training loss, the updated parameters, and the updated optimizer state.
"""
# Build the position and attention mask vectors.
positions, attention_mask = get_attention_mask_and_positions(example.input_tokens, pad_id)
# The forward and backward passes.
train_loss, grads = jax.value_and_grad(forward_and_loss_fn)(params,
model=model,
input_tokens=example.input_tokens,
input_mask=example.target_mask,
positions=positions,
attention_mask=attention_mask)
# Update the parameters.
updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)
return train_loss, params, opt_state
Geriye doğru geçiş olmadan validation_step işlevini derleyin:
def validation_step(model: transformer_lib.Transformer,
params,
pad_id: int,
example: TrainingInput,
):
positions, attention_mask = get_attention_mask_and_positions(example.input_tokens, pad_id)
val_loss = forward_and_loss_fn(params,
model=model,
input_tokens=example.input_tokens,
input_mask=example.target_mask,
positions=positions,
attention_mask=attention_mask)
return val_loss
SGD optimize edici için optax.sgd kullanarak eğitim döngüsünü tanımlayın:
@chex.dataclass(frozen=True)
class TrainingConfig:
learning_rate: float
num_epochs: int
eval_every_n: int
batch_size: int
max_steps: int | None = None
def train_loop(
model: transformer_lib.Transformer,
params,
dataset_builder: MTNTDatasetBuilder,
training_cfg: TrainingConfig):
# Apply `jax.jit` on the training step, making the whole loop much more efficient.
compiled_train_step = jax.jit(train_step, static_argnames=['model', 'optimizer'])
# Apply `jax.jit` on the validation step.
compiled_validation_step = jax.jit(validation_step, static_argnames=['model'])
# To save memory, use the SGD optimizer instead of the usual Adam optimizer.
# Note that for this specific example, SGD is more than enough.
optimizer = optax.sgd(training_cfg.learning_rate)
opt_state = optimizer.init(params)
# Build the training dataset.
train_ds = dataset_builder.get_train_dataset(batch_size=training_cfg.batch_size,
num_epochs=training_cfg.num_epochs)
train_ds = train_ds.as_numpy_iterator()
# Build the validation dataset, with a limited number of samples for this demo.
validation_ds = dataset_builder.get_validation_dataset(batch_size=training_cfg.batch_size)
validation_ds = validation_ds.take(50)
n_steps = 0
avg_loss=0
# A first round of the validation loss.
n_steps_eval = 0
eval_loss = 0
val_iterator = validation_ds.as_numpy_iterator()
for val_example in val_iterator:
eval_loss += compiled_validation_step(model,
params,
dataset_builder._tokenizer.pad_id,
val_example)
n_steps_eval += 1
print(f"Start, validation loss: {eval_loss/n_steps_eval}")
for train_example in train_ds:
train_loss, params, opt_state = compiled_train_step(model=model,
params=params,
optimizer=optimizer,
opt_state=opt_state,
pad_id=dataset_builder._tokenizer.pad_id,
example=train_example)
n_steps += 1
avg_loss += train_loss
if n_steps % training_cfg.eval_every_n == 0:
eval_loss = 0
n_steps_eval = 0
val_iterator = validation_ds.as_numpy_iterator()
for val_example in val_iterator:
eval_loss += compiled_validation_step(model,
params,
dataset_builder._tokenizer.pad_id,
val_example)
n_steps_eval +=1
avg_loss /= training_cfg.eval_every_n
eval_loss /= n_steps_eval
print(f"STEP {n_steps} training loss: {avg_loss} - eval loss: {eval_loss}")
avg_loss=0
if training_cfg.max_steps is not None and n_steps > training_cfg.max_steps:
break
return params
Hafızaya sığması için Gemma modelini sınırlı sayıda adımda (SEQ_SIZE) ince ayar yapmaya başlayın:
SEQ_SIZE = 25
tokenizer = GemmaTokenizer(vocab)
dataset_builder= MTNTDatasetBuilder(tokenizer, SEQ_SIZE)
training_cfg = TrainingConfig(learning_rate=1e-4,
num_epochs=1,
eval_every_n=20,
batch_size=1,
max_steps=100)
params = train_loop(model=model_2b,
params={'params': params['transformer']},
dataset_builder=dataset_builder,
training_cfg=training_cfg)
Start, validation loss: 10.647212982177734 STEP 20 training loss: 3.3015992641448975 - eval loss: 2.686880111694336 STEP 40 training loss: 5.375057220458984 - eval loss: 2.6751961708068848 STEP 60 training loss: 2.6599338054656982 - eval loss: 2.663877010345459 STEP 80 training loss: 4.822389125823975 - eval loss: 2.3333375453948975 STEP 100 training loss: 2.0131142139434814 - eval loss: 2.360811948776245
Her adım sayısında hem eğitim kaybı hem de doğrulama kaybının azalması gerekiyordu.
gemma.sampler.Sampler ile bir sampler oluşturun. Gemma modeli kontrol noktasını ve jeton belirleyiciyi kullanır.
sampler = sampler_lib.Sampler(
transformer=model_2b,
vocab=vocab,
params=params['params'],
)
Modelinizin çeviri yapıp yapamayacağını kontrol etmek için sampler özelliğini kullanın. gemma.sampler.Sampler içindeki total_generation_steps bağımsız değişkeni, yanıt oluşturulurken gerçekleştirilen adımların sayısıdır. Girişin eğitim biçimiyle eşleştiğinden emin olmak için sonunda yeni satır karakteri bulunan Translate this into French:\n ön ekini kullanın. Bu, modelin çeviriye başlaması için sinyal gönderir.
sampler(
["Translate this into French:\nHello, my name is Morgane.\n"],
total_generation_steps=100,
).text
["C'est Bonjour, mon nom est Morgane.C'est Bonjour, mon nom est Morgane."]
Daha fazla bilgi
- Bu eğitimde kullandığınız
gemma.params,gemma.transformervegemma.samplergibi modüllerin doküman dizelerini içeren Google DeepMindgemmakitaplığını GitHub'da bulabilirsiniz. - Şu kitaplıkların kendi doküman siteleri vardır: core JAX, Flax, Chex, Optax ve Orbax.
sentencepiecetokenizleyici/detokenizer belgeleri için Google'ınsentencepieceGitHub deposuna göz atın.kagglehubdokümanları için Kaggle'ınkagglehubGitHub deposundakiREADME.mdkoleksiyonuna göz atın.- Gemma modellerini Google Cloud Vertex AI ile nasıl kullanacağınızı öğrenin.
- Google Cloud TPU'ları (v3-8 ve daha yeni bir sürüm) kullanıyorsanız en son
jax[tpu]paketine (!pip install -U jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html) güncellediğinizden emin olun, çalışma zamanını yeniden başlatın vejaxilejaxlibsürümlerinin eşleşip eşleşmediğini (!pip list | grep jax) kontrol edin. Bu,jaxlibvejaxsürümlerinin uyuşmaması nedeniyle oluşabilecekRuntimeErrorhatalarını önleyebilir. Diğer JAX yükleme talimatları için JAX dokümanlarına bakın.