|
from __gin__ import dynamic_registration |
|
import tasks |
|
import seqio |
|
import optax |
|
|
|
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 = %gin.REQUIRED |
|
TASK_FEATURE_LENGTHS = {"inputs": 256, "targets": 256} |
|
INITIAL_CHECKPOINT_PATH = %gin.REQUIRED |
|
LR = %gin.REQUIRED |
|
TRAIN_STEPS = %gin.REQUIRED |
|
USE_CACHED_TASKS = False |
|
DROPOUT_RATE = 0.1 |
|
RANDOM_SEED = 0 |
|
BATCH_SIZE = 32 |
|
EVAL_PERIOD = 200 |
|
|
|
|
|
infer_eval/utils.DatasetConfig: |
|
task_feature_lengths = %TASK_FEATURE_LENGTHS |
|
|
|
|
|
utils.SaveCheckpointConfig: |
|
period = 200 |
|
keep = 1 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
utils.create_learning_rate_scheduler: |
|
factors = 'constant * rsqrt_decay' |
|
base_learning_rate = %LR |
|
warmup_steps = 187 |
|
|
|
|
|
|
|
|
|
|
|
|