|
from __gin__ import dynamic_registration |
|
import __main__ as eval_script |
|
import seqio |
|
from t5.data import mixtures |
|
from t5x import adafactor |
|
from t5x.examples.t5 import network |
|
from t5x import models |
|
from t5x import partitioning |
|
from t5x import utils |
|
import tasks |
|
|
|
|
|
|
|
CHECKPOINT_PATH = \ |
|
'gs://nb-t5x-us-central2/finetuned/v2_norwegian_NCC_plus_English_t5x_base_1_500_000_sentiment/checkpoint_1510000' |
|
DROPOUT_RATE = 0.0 |
|
EVAL_OUTPUT_DIR = './log/' |
|
LABEL_SMOOTHING = 0.0 |
|
LOSS_NORMALIZING_FACTOR = None |
|
MIXTURE_OR_TASK_NAME = 'sentiment' |
|
MODEL = @models.EncoderDecoderModel() |
|
OPTIMIZER = @adafactor.Adafactor() |
|
SPLIT = 'test' |
|
VOCABULARY = @seqio.SentencePieceVocabulary() |
|
Z_LOSS = 0.0001 |
|
|
|
|
|
|
|
adafactor.Adafactor.decay_rate = 0.8 |
|
adafactor.Adafactor.logical_factor_rules = \ |
|
@adafactor.standard_logical_factor_rules() |
|
adafactor.Adafactor.step_offset = 0 |
|
|
|
|
|
|
|
utils.DatasetConfig.batch_size = 16 |
|
utils.DatasetConfig.mixture_or_task_name = %MIXTURE_OR_TASK_NAME |
|
utils.DatasetConfig.seed = 42 |
|
utils.DatasetConfig.shuffle = False |
|
utils.DatasetConfig.split = %SPLIT |
|
utils.DatasetConfig.task_feature_lengths = {'inputs': 512, 'targets': 2} |
|
|
|
|
|
|
|
models.EncoderDecoderModel.input_vocabulary = %VOCABULARY |
|
models.EncoderDecoderModel.label_smoothing = %LABEL_SMOOTHING |
|
models.EncoderDecoderModel.loss_normalizing_factor = %LOSS_NORMALIZING_FACTOR |
|
models.EncoderDecoderModel.module = @network.Transformer() |
|
models.EncoderDecoderModel.optimizer_def = %OPTIMIZER |
|
models.EncoderDecoderModel.output_vocabulary = %VOCABULARY |
|
models.EncoderDecoderModel.z_loss = %Z_LOSS |
|
|
|
|
|
|
|
eval_script.evaluate.dataset_cfg = @utils.DatasetConfig() |
|
eval_script.evaluate.model = %MODEL |
|
eval_script.evaluate.output_dir = %EVAL_OUTPUT_DIR |
|
eval_script.evaluate.partitioner = @partitioning.PjitPartitioner() |
|
eval_script.evaluate.restore_checkpoint_cfg = @utils.RestoreCheckpointConfig() |
|
|
|
|
|
|
|
partitioning.PjitPartitioner.num_partitions = 2 |
|
|
|
|
|
|
|
utils.RestoreCheckpointConfig.mode = 'specific' |
|
utils.RestoreCheckpointConfig.path = %CHECKPOINT_PATH |
|
|
|
|
|
|
|
seqio.SentencePieceVocabulary.sentencepiece_model_file = \ |
|
'gs://t5-data/vocabs/mc4.250000.100extra/sentencepiece.model' |
|
|
|
|
|
|
|
network.T5Config.dropout_rate = %DROPOUT_RATE |
|
network.T5Config.dtype = 'bfloat16' |
|
network.T5Config.emb_dim = 768 |
|
network.T5Config.head_dim = 64 |
|
network.T5Config.logits_via_embedding = False |
|
network.T5Config.mlp_activations = ('gelu', 'linear') |
|
network.T5Config.mlp_dim = 2048 |
|
network.T5Config.num_decoder_layers = 12 |
|
network.T5Config.num_encoder_layers = 12 |
|
network.T5Config.num_heads = 12 |
|
network.T5Config.vocab_size = 250112 |
|
|
|
|
|
|
|
network.Transformer.config = @network.T5Config() |
|
|
