finetune_mt5_large_lm.gin

from __gin__ import dynamic_registration

include 't5x/examples/t5/mt5/large.gin'
include 't5x/configs/runs/finetune.gin'

# Register necessary SeqIO Tasks/Mixtures.
import t5.data.mixtures

import tasks

# INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/mt5_large/checkpoint_1000000"
INITIAL_CHECKPOINT_PATH = "gs://bertin-project/t5/mt5_large/checkpoint_4720000"
MIXTURE_OR_TASK_NAME = "mc4_es_gaussian_prefix_lm_mt5"
TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 512}
TRAIN_STEPS = 5000000  # Original 100000, we start from 4720000, so let's do just almost 4 times that ((2048/(64*8))=4)

DROPOUT_RATE = 0.0
BATCH_SIZE = 56  # It is 64 per TPU