new script for partitioning. Based on base_wmt14enfr

finetune_large_mt5_sentencefix_v4_16.gin

@@ -25,15 +25,35 @@ RANDOM_SEED = 0
 INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/mt5_large/checkpoint_1000000"
 
 train_script.train:
-  eval_period = 500
+  eval_period = 100
+
+utils.RestoreCheckpointConfig:
+  path = %INITIAL_CHECKPOINT_PATH
+  mode = 'specific'
+
+train_script.train:
+  train_dataset_cfg = @train/utils.DatasetConfig()
+  train_eval_dataset_cfg = @train_eval/utils.DatasetConfig()
+  infer_eval_dataset_cfg = @infer_eval/utils.DatasetConfig()
 
-# `num_decodes` is equivalent to a beam size in a beam search decoding.
 models.EncoderDecoderModel.predict_batch_with_aux.num_decodes = 4
 
-partitioning.PjitPartitioner.num_partitions = 4
+infer_eval/utils.DatasetConfig:
+  mixture_or_task_name = %MIXTURE_OR_TASK_NAME
+  task_feature_lengths = %TASK_FEATURE_LENGTHS
+  split = 'validation'
+  batch_size = 64
+  shuffle = False
+  seed = 42
+  use_cached = %USE_CACHED_TASKS
+  pack = False
+  module = %MIXTURE_OR_TASK_MODULE
 
-#from t5.models import mesh_transformer
-#import t5.models
-#mesh_transformer.learning_rate_schedules.constant_learning_rate.learning_rate = 0.0005
-#run.learning_rate_schedule = @learning_rate_schedules.constant_learning_rate
+seqio.Evaluator:
+  logger_cls = [@seqio.PyLoggingLogger, @seqio.TensorBoardLogger, @seqio.JSONLogger]
+  num_examples = None  # Use all examples in the dataset.
+  use_memory_cache = True
+
+
+partitioning.PjitPartitioner.num_partitions = 4