small_pretrain.gin

# Register necessary SeqIO Tasks/Mixtures.
from __gin__ import dynamic_registration
from t5x import utils
import ul2_tasks
import __main__ as train_script

include 'small.gin'
include 't5x/configs/runs/pretrain.gin'


# ------------------- Training specification overrides --------------------------
train_script.train:
  eval_period = 2000
  stats_period = 100

utils.SaveCheckpointConfig:
    period = 50000
    keep = 4


MIXTURE_OR_TASK_NAME = "pretrain_biological_ul2"
USE_CACHED_TASKS = False
TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 512}
TRAIN_STEPS = 500000
DROPOUT_RATE = 0.0
BATCH_SIZE = 256