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 = 32
batch_fun.eval_batch_size = 32

# Parameters for inputs:
# ==============================================================================
inputs.dataset_name = 't2t_algorithmic_identity_vocab95_train20_eval30'

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

# Parameters for preprocess_fun:
# ==============================================================================
shuffle_and_batch_data.preprocess_fun=@trax.inputs.concat_preprocess

# Parameters for MultifactorSchedule:
# ==============================================================================
MultifactorSchedule.constant = 0.05
MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
MultifactorSchedule.warmup_steps = 16000

# Parameters for PositionLookupTransformerLM:
# ==============================================================================
PositionLookupTransformerLM.d_feature = 256
PositionLookupTransformerLM.d_feedforward = 512
PositionLookupTransformerLM.dropout = 0.01
PositionLookupTransformerLM.max_len = 100
PositionLookupTransformerLM.n_heads = 4
PositionLookupTransformerLM.n_layers = 3
PositionLookupTransformerLM.vocab_size = 128

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