import tensor2tensor.trax.inputs
import tensor2tensor.trax.models
import tensor2tensor.trax.optimizers
import tensor2tensor.trax.trax

# Parameters for batch_fun:
# ==============================================================================
batch_fun.batch_size_per_device = 128
batch_fun.eval_batch_size = 128
batch_fun.max_eval_length = 2048

# Parameters for inputs:
# ==============================================================================
inputs.data_dir = None
inputs.dataset_name = 't2t_languagemodel_lm1b32k'
inputs.input_name = 'targets'

# Parameters for mask:
# ==============================================================================
masked_mean.mask_id = 0

# Parameters for MultifactorSchedule:
# ==============================================================================
MultifactorSchedule.constant = 0.1
MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
MultifactorSchedule.warmup_steps = 8000

# Parameters for preprocess_fun:
# ==============================================================================
shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
lm1b_preprocess.max_target_length = 512
lm1b_preprocess.max_eval_target_length = 2048

# Parameters for train:
# ==============================================================================
train.eval_frequency = 1000
train.eval_steps = 10
train.inputs = @trax.inputs.inputs
train.model = @trax.models.TransformerLM
train.optimizer = @trax.optimizers.Adam
train.train_steps = 100000

# Parameters for TransformerLM:
# ==============================================================================
TransformerLM.d_feature = 512
TransformerLM.d_feedforward = 2048
TransformerLM.dropout = 0.1
TransformerLM.max_len = 2048
TransformerLM.mode = 'train'
TransformerLM.n_heads = 8
TransformerLM.n_layers = 6
TransformerLM.vocab_size = 32000
