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

# Parameters for AudioDataModule:
# ==============================================================================
AudioDataModule.num_workers = 0

# Parameters for AudioDataset:
# ==============================================================================
AudioDataset.half_precision = True
AudioDataset.mono = True
AudioDataset.new_freq = 24000
AudioDataset.num_frames = 480000
AudioDataset.orig_freq = 16000

# Parameters for build_dev_datamodule:
# ==============================================================================
build_dev_datamodule.datamodule = @discotube_text_audio

# Parameters for build_module:
# ==============================================================================
build_module.ckpt_path = \
    '/gpfs/projects/upf97/logs/mtg-text-audio/bayc7zny/checkpoints/epoch=14-step=305670.ckpt'
build_module.module = @modules.maskingmodel.MaskingModel
build_module.net = @nets.conformer.Conformer
build_module.representation = \
    [@nets.melspectrogram.MelSpectrogram,
     @nets.waveform.Waveform,
     @nets.cqt.CQT,
     @nets.encodec.EnCodec]


def_module.module = @clap

CLAP.audio_encoder = @OMARRQ
CLAP.text_encoder = @AllMPNetBaseV2
CLAP.proj_size = 512
CLAP.temp = 0.1
CLAP.lr = 1e-4
CLAP.weight_decay = 1e-2
CLAP.seed = 0
CLAP.train_audio_encoder = True
CLAP.train_text_encoder = False
CLAP.tokenizers_parallelism = False

OMARRQ.model_id = "mtg-upf/clap_omarrq_mp_small_music"
AllMPNetBaseV2.model_id = "sentence-transformers/all-mpnet-base-v2"

# Parameters for Conformer:
# ==============================================================================
Conformer.alpha_deepnorm = 2.21
Conformer.beta_deepnorm = 0.0026
Conformer.conv_kernel_size = 5
Conformer.depth = 12
Conformer.dropout = 0.2
Conformer.embed_dim = 512
Conformer.input_dropout = 0.0
Conformer.mlp_ratio = 4.0
Conformer.mlp_residual_factor = 4.0
Conformer.num_heads = 8
Conformer.num_patches = None
Conformer.use_deepnorm = True
Conformer.use_rope = True

# Parameters for CQT:
# ==============================================================================
CQT.bins_per_octave = 24
CQT.f_min = 32.703
CQT.hop_len = 320
CQT.logC = True
CQT.magnitude = True
CQT.n_bins = 188
CQT.norm_mean = 4.754879065310596
CQT.norm_std = 1.9055732535255916
CQT.patch_size = (188, 4)
CQT.power = 2
CQT.sr = %new_freq


# Parameters for EnCodec:
# ==============================================================================
EnCodec.norm_type = 'global'
EnCodec.orig_sr = %new_freq
EnCodec.patch_size = (128, 4)
EnCodec.stats_path = None
EnCodec.weights_path = '/gpfs/scratch/upf97/model_weights/encodec_24khz/'

# Parameters for MaskingModel:
# ==============================================================================
MaskingModel.codebook_dim = 16
MaskingModel.codebook_size = 8196
MaskingModel.diff_input = False
MaskingModel.input_representation = @nets.waveform.Waveform
MaskingModel.lr = 0.0001
MaskingModel.mask_prob = 0.6
MaskingModel.mask_seconds = 0.4
MaskingModel.num_codebooks = 1
MaskingModel.plot_tokens = False
MaskingModel.seed = 0
MaskingModel.weight_decay = 0.01

# Parameters for MelSpectrogram:
# ==============================================================================
MelSpectrogram.freq_mask_param = 0
MelSpectrogram.hop_len = 320
MelSpectrogram.mel_scale = 'slaney'
MelSpectrogram.n_mel = 96
MelSpectrogram.norm = 'slaney'
MelSpectrogram.norm_mean = 2.06755686098554
MelSpectrogram.norm_std = 1.268292820667291
MelSpectrogram.patch_size = (96, 4)
MelSpectrogram.power = 2
MelSpectrogram.sr = %new_freq
MelSpectrogram.stretch_factor = 1
MelSpectrogram.time_mask_param = 0
MelSpectrogram.win_len = 512

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