Visualizza su ai.google.dev | Esegui in Google Colab | Apri in Vertex AI | Visualizza il codice sorgente su GitHub |
Panoramica
Gemma è una famiglia di modelli linguistici di grandi dimensioni aperti, leggeri e all'avanguardia, basati sulla ricerca e sulla tecnologia di Google DeepMind Gemini. Questo tutorial mostra come perfezionare il modello Gemma 2B Instruct per un'attività di traduzione inglese-francese utilizzando la libreria gemma
di Google DeepMind, JAX (una libreria di calcolo numerico ad alte prestazioni), Flax (la libreria di rete neurale basata su JAX), Chex (una libreria di utilità per la scrittura di set di dati JAX affidabili per l'elaborazione di testi JAX e JAX), Optax e Text2 (Text2) Sebbene Lino non sia utilizzato direttamente in questo blocco note, è stato utilizzato per creare Gemma.
La libreria gemma
è stata scritta con JAX, Flax, Orbax (una libreria basata su JAX per l'addestramento di utilità, come il checkpointing) e SentencePiece (una libreria tokenizzatore/detokenizzatore).
Configurazione
1. Configurare l'accesso di Kaggle per Gemma
Per completare questo tutorial, devi prima seguire le istruzioni di configurazione di Gemma che mostrano come effettuare le seguenti operazioni:
- Accedi a Gemma su kaggle.com.
- Seleziona un runtime Colab con risorse sufficienti per eseguire il modello Gemma.
- Genera e configura un nome utente e una chiave API di Kaggle.
Dopo aver completato la configurazione di Gemma, passa alla sezione successiva, dove potrai impostare le variabili di ambiente per l'ambiente Colab.
2. Imposta le variabili di ambiente
Imposta le variabili di ambiente per KAGGLE_USERNAME
e KAGGLE_KEY
. Quando viene visualizzato il messaggio "Concedi l'accesso?", accetta di fornire l'accesso segreto.
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. Installa la libreria gemma
Al momento, l'accelerazione hardware di Colab senza costi è sufficiente per eseguire questo blocco note. Se usi Colab Pay As You Go o Colab Pro, fai clic su Modifica > Impostazioni blocco note > Seleziona GPU A100 > Salva per attivare l'accelerazione hardware.
Dopodiché devi installare la libreria Google DeepMind gemma
da github.com/google-deepmind/gemma
. Se viene visualizzato un errore relativo al " resolver di dipendenze di pip", generalmente puoi ignorarlo.
pip install -q git+https://github.com/google-deepmind/gemma.git
4. Importa librerie
Questo blocco note utilizza Flax (per le reti neurali), il core JAX, SentencePiece (per la tokenizzazione), Chex (una libreria di utilità per scrivere codice JAX affidabile) e set di dati TensorFlow.
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
Carica il modello Gemma
Carica il modello Gemma con kagglehub.model_download
, che accetta tre argomenti:
handle
: l'handle del modello di Kagglepath
: (stringa facoltativa) il percorso localeforce_download
: (booleano facoltativo) costringe a scaricare di nuovo il modello
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
Controlla la posizione delle ponderazioni del modello e del tokenizzatore, quindi imposta le variabili del percorso. La directory del tokenizzatore si trova nella directory principale in cui hai scaricato il modello, mentre i pesi del modello si trovano in una sottodirectory. Ad esempio:
- Il file
tokenizer.model
si troverà in/LOCAL/PATH/TO/gemma/flax/2b-it/2
). - Il checkpoint del modello si trova in
/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
Carica e prepara il set di dati MTNT e il tokenizzatore Gemma
Utilizzerai il set di dati MTNT (Machine Translation of Noisy Text), disponibile nei set di dati TensorFlow.
Scarica la parte del set di dati dall'inglese al francese del set di dati MTNT, quindi campiona due esempi. Ogni campione nel set di dati contiene due voci: src
, la frase originale in inglese, e dst
, la traduzione in francese corrispondente.
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?'
Carica il tokenizzatore Gemma, creato utilizzando sentencepiece.SentencePieceProcessor
:
vocab = spm.SentencePieceProcessor()
vocab.Load(TOKENIZER_PATH)
True
Personalizza la SentencePieceProcessor
per l'attività di traduzione dall'inglese al francese. Poiché ottimizzerai la parte inglese del modello Gemma, dovrai apportare alcune modifiche, ad esempio:
Prefisso di input: l'aggiunta di un prefisso comune a ogni input segnala l'attività di traduzione. Ad esempio, potresti utilizzare un prompt con un prefisso come
Translate this into French: [INPUT_SENTENCE]
.Il suffisso di inizio della traduzione: l'aggiunta di un suffisso alla fine di ogni prompt indica al modello Gemma esattamente quando iniziare il processo di traduzione. Una nuova riga dovrebbe farlo.
Token del modello linguistico: i modelli Gemma si aspettano un token di "inizio sequenza" all'inizio di ogni sequenza, quindi aggiungere un token di "fine sequenza" alla fine di ogni esempio di addestramento dovrebbe essere sufficiente.
Crea un wrapper personalizzato intorno a
SentencePieceProcessor
nel seguente modo:
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())
Prova creando un'istanza del nuovo GemmaTokenizer
personalizzato e poi applicandolo a un piccolo campione del set di dati MTNT:
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]
Crea un caricatore di dati per l'intero set di dati MTNT:
@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
Prova MTNTDatasetBuilder
creando di nuovo un'istanza del GemmaTokenizer
personalizzato, quindi applicandolo al set di dati MTNT ed eseguendo il campionamento di due esempi:
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]]
Configura il modello
Prima di iniziare il perfezionamento del modello Gemma, devi configurarlo.
Innanzitutto, carica e formatta il checkpoint del modello Gemma con il metodo gemma.params.load_and_format_params
:
params = params_lib.load_and_format_params(CKPT_PATH)
Per caricare automaticamente la configurazione corretta dal checkpoint del modello Gemma, utilizza gemma.transformer.TransformerConfig
. L'argomento cache_size
è il numero di passaggi temporali nella cache Transformer
di Gemma. Successivamente, crea l'istanza del modello Gemma come model_2b
con gemma.transformer.Transformer
(che eredita da flax.linen.Module
).
config_2b = transformer_lib.TransformerConfig.from_params(
params,
cache_size=30
)
model_2b = transformer_lib.Transformer(config=config_2b)
Ottimizza il modello
In questa sezione imparerai a:
- Usa la classe
gemma.transformer.Transformer
per creare la funzione pass and loss. - Crea i vettori della maschera di attenzione e della posizione per i token
- Crea una funzione per i passaggi di addestramento con Flax.
- Crea il passaggio di convalida senza il passaggio a ritroso.
- Crea il ciclo di addestramento.
- Ottimizzare il modello Gemma.
Definisci il passaggio in avanti e la funzione loss utilizzando la classe gemma.transformer.Transformer
. L'Transformer
di Gemma eredita da flax.linen.Module
e offre due metodi essenziali:
init
: inizializza i parametri del modello.apply
: esegue la funzione__call__
del modello utilizzando un determinato insieme di parametri.Poiché stai utilizzando le ponderazioni Gemma preaddestrate, non è necessario utilizzare la funzione
init
.
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
La classe gemma.transformer.Transformer
richiede un vettore attention_mask
e positions
insieme a ogni input. Puoi generarle creando una funzione personalizzata che utilizza Transformer.build_positions_from_mask
e Transformer.make_causal_attn_mask
:
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
Crea la funzione train_step
che esegue il passaggio a ritroso e aggiorna i parametri del modello di conseguenza, dove:
jax.value_and_grad
consente di valutare la funzione di perdita e i gradienti durante i passaggi in avanti e indietro.optax.apply_updates
consente di aggiornare i parametri.
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
Crea la funzione validation_step
senza il passaggio a ritroso:
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
Definisci il ciclo di addestramento utilizzando optax.sgd
per l'ottimizzatore SGD:
@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
Inizia a perfezionare il modello Gemma eseguendo un numero limitato di passaggi (SEQ_SIZE
) per assicurarti che rientri nella memoria:
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
Sia la perdita per l'addestramento che la perdita per la convalida dovrebbero essere diminuite a ogni numero di passaggi.
Crea un sampler
con gemma.sampler.Sampler
. Utilizza il checkpoint del modello Gemma e il tokenizzatore.
sampler = sampler_lib.Sampler(
transformer=model_2b,
vocab=vocab,
params=params['params'],
)
Utilizza sampler
per verificare se il modello può eseguire la traduzione. L'argomento total_generation_steps
in gemma.sampler.Sampler
indica il numero di passaggi eseguiti durante la generazione di una risposta. Per assicurarti che l'input corrisponda al formato di addestramento, utilizza il prefisso Translate this into French:\n
con un carattere di nuova riga alla fine. In questo modo, il modello può iniziare la traduzione.
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."]
Scopri di più
- Puoi scoprire di più sulla libreria Google DeepMind
gemma
su GitHub, che contiene i docstring dei moduli utilizzati in questo tutorial, comegemma.params
,gemma.transformer
egemma.sampler
. - Le seguenti librerie hanno i propri siti di documentazione: core JAX, Flax, Chex, Optax e Orbax.
- Per la documentazione relativa ai tokenizzatori/detokenizzatori di
sentencepiece
, consulta il repository GitHub disentencepiece
di Google. - Per la documentazione relativa a
kagglehub
, dai un'occhiata aREADME.md
sul repository GitHub di Kaggle,kagglehub
. - Scopri come utilizzare i modelli Gemma con Vertex AI di Google Cloud.
- Se utilizzi Google Cloud TPU (v3-8 e versioni successive), assicurati di eseguire anche l'aggiornamento al pacchetto
jax[tpu]
più recente (!pip install -U jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
), riavvia il runtime e verifica che le versionijax
ejaxlib
corrispondano (!pip list | grep jax
). Ciò può evitare che si possa verificareRuntimeError
a causa della mancata corrispondenza delle versionijaxlib
ejax
. Per ulteriori istruzioni sull'installazione di JAX, consulta la documentazione di JAX.