|
from __gin__ import dynamic_registration |
|
import tasks |
|
|
|
import __main__ as train_script |
|
from t5.data import mixtures |
|
from t5x import models |
|
from t5x import partitioning |
|
from t5x import utils |
|
|
|
include "t5x/examples/t5/mt5/base.gin" |
|
include "t5x/configs/runs/finetune.gin" |
|
|
|
MIXTURE_OR_TASK_NAME = "categorise" |
|
TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 2} |
|
TRAIN_STEPS = 1_510_000 # 1000000 pre-trained steps + 10000 fine-tuning steps. |
|
USE_CACHED_TASKS = False |
|
DROPOUT_RATE = 0.1 |
|
RANDOM_SEED = 0 |
|
|
|
#Fixing a small error |
|
infer_eval/utils.DatasetConfig: |
|
task_feature_lengths = %TASK_FEATURE_LENGTHS |
|
|
|
# Pere: Only necessary if we load a t5 model. We can start with an t5x model here |
|
# `LOSS_NORMALIZING_FACTOR`: When fine-tuning a model that was pre-trained |
|
# using Mesh Tensorflow (e.g. the public T5 / mT5 / ByT5 models), this should be |
|
# set to `pretraining batch_size` * `target_token_length`. For T5 and T5.1.1: |
|
# `2048 * 114`. For mT5: `1024 * 229`. For ByT5: `1024 * 189`. |
|
#LOSS_NORMALIZING_FACTOR = 234496 |
|
|
|
#INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/mt5_base/checkpoint_1000000" |
|
#INITIAL_CHECKPOINT_PATH = "gs://nb-t5x-us-central2/pk_nb_t5x_base_run1/checkpoint_1100000" |
|
#INITIAL_CHECKPOINT_PATH = "gs://nb-t5x-us-central2/norwegian_t5x_base/checkpoint_1360000" |
|
#INITIAL_CHECKPOINT_PATH = "gs://nb-t5x-us-central2/pk_nb_t5x_base_run1_lr_1/checkpoint_1100000" |
|
#INITIAL_CHECKPOINT_PATH = "gs://nb-t5x-us-central2/pk_nb_t5x_base_scandinavian/checkpoint_1100000" |
|
#INITIAL_CHECKPOINT_PATH = "gs://nb-t5x-us-central2/norwegian_t5x_base/checkpoint_2000000" |
|
INITIAL_CHECKPOINT_PATH = "gs://nb-t5x-us-central2/norwegian_NCC_plus_English_t5x_base/checkpoint_1500000" |
|
|
|
|
|
|
|
#train_script.train: |
|
# eval_period = 500 |
|
# partitioner = @partitioning.ModelBasedPjitPartitioner() |
|
# partitioning.PjitPartitioner.num_partitions = 1 |
|
|
|
# `num_decodes` is equivalent to a beam size in a beam search decoding. |
|
# models.EncoderDecoderModel.predict_batch_with_aux.num_decodes = 1 |
|
|
|
#mesh_transformer.learning_rate_schedules.constant_learning_rate.learning_rate = 0.0005 |
|
#run.learning_rate_schedule = @learning_rate_schedules.constant_learning_rate |
|
|
|
|