# Parameters for build_module:
# ==============================================================================
omar_rq.utils.build_module.ckpt_path = 'model.ckpt'
omar_rq.utils.build_module.module = @omar_rq.modules.maskingmodel.MaskingModel
omar_rq.utils.build_module.net = @omar_rq.nets.conformer.Conformer
omar_rq.utils.build_module.representation = \
    [@omar_rq.nets.cqt.CQT,
     @omar_rq.nets.encodec.EnCodec,
     @omar_rq.nets.melspectrogram.MelSpectrogram,
     @omar_rq.nets.waveform.Waveform]

# Macros:
# ==============================================================================
new_freq = 24000

# Parameters for omar_rq.nets.conformer.Conformer:
# ==============================================================================
omar_rq.nets.conformer.Conformer.alpha_deepnorm = 2.6321480259049848
omar_rq.nets.conformer.Conformer.beta_deepnorm = 0.022386873579657126
omar_rq.nets.conformer.Conformer.conv_kernel_size = 5
omar_rq.nets.conformer.Conformer.depth = 24
omar_rq.nets.conformer.Conformer.dropout = 0.2
omar_rq.nets.conformer.Conformer.embed_dim = 1024
omar_rq.nets.conformer.Conformer.input_dropout = 0.0
omar_rq.nets.conformer.Conformer.mlp_ratio = 4.0
omar_rq.nets.conformer.Conformer.mlp_residual_factor = 4.0
omar_rq.nets.conformer.Conformer.num_heads = 8
omar_rq.nets.conformer.Conformer.num_patches = None
omar_rq.nets.conformer.Conformer.patch_size = None
omar_rq.nets.conformer.Conformer.use_deepnorm = True
omar_rq.nets.conformer.Conformer.use_rope = True

# Parameters for omar_rq.modules.maskingmodel.MaskingModel:
# ==============================================================================
omar_rq.modules.maskingmodel.MaskingModel.codebook_dim = 1
omar_rq.modules.maskingmodel.MaskingModel.codebook_size = 7776
omar_rq.modules.maskingmodel.MaskingModel.diff_input = False
omar_rq.modules.maskingmodel.MaskingModel.input_representation = @nets.waveform.Waveform
omar_rq.modules.maskingmodel.MaskingModel.lr = 0.0001
omar_rq.modules.maskingmodel.MaskingModel.mask_prob = 0.6
omar_rq.modules.maskingmodel.MaskingModel.mask_seconds = 0.4
omar_rq.modules.maskingmodel.MaskingModel.masking_noise_type = 'random_normal'
omar_rq.modules.maskingmodel.MaskingModel.num_codebooks = 1
omar_rq.modules.maskingmodel.MaskingModel.plot_tokens = False
omar_rq.modules.maskingmodel.MaskingModel.quantizer_type = 'finite_scalar_quantizer'
omar_rq.modules.maskingmodel.MaskingModel.seed = 0
omar_rq.modules.maskingmodel.MaskingModel.weight_decay = 0.01

# Parameters for Waveform:
# ==============================================================================
Waveform.norm_mean = None
Waveform.norm_std = None
Waveform.patch_size = (1, 960)
Waveform.sr = %new_freq