# ############################################################################
# Model: WAV2VEC XLSR model for Accent Recognition (Spanish)
# see paper: https://arxiv.org/abs/2305.18283
# ############################################################################

# Hparams NEEDED
HPARAMS_NEEDED: ["encoder_dim", "out_n_neurons", "label_encoder", "softmax"]
# Modules Needed
MODULES_NEEDED: ["wav2vec2", "avg_pool", "output_mlp"]

# Feature parameters
# wav2vec2_hub: facebook/wav2vec2-base
wav2vec2_hub: "facebook/wav2vec2-large-xlsr-53"

# Pretrain folder (HuggingFace)
pretrained_path: Jzuluaga/accent-id-commonaccent_xlsr-es-spanish
# URL for the biggest Fairseq english wav2vec2 model.

# parameters
encoder_dim: 1024
out_n_neurons: 6

wav2vec2: !new:speechbrain.lobes.models.huggingface_wav2vec.HuggingFaceWav2Vec2
    source: !ref <wav2vec2_hub>
    output_norm: True
    freeze: True
    save_path: wav2vec2_checkpoints

# Mean and std normalization of the input features
mean_var_norm_emb: !new:speechbrain.processing.features.InputNormalization
    norm_type: sentence
    std_norm: False

avg_pool: !new:speechbrain.nnet.pooling.StatisticsPooling
    return_std: False

output_mlp: !new:speechbrain.nnet.linear.Linear
    input_size: !ref <encoder_dim>
    n_neurons: !ref <out_n_neurons>
    bias: False

model: !new:torch.nn.ModuleList
    - [!ref <output_mlp>]

modules:
    mean_var_norm_emb: !ref <mean_var_norm_emb>
    wav2vec2: !ref <wav2vec2>
    output_mlp: !ref <output_mlp>
    avg_pool: !ref <avg_pool>

softmax: !new:speechbrain.nnet.activations.Softmax
        
label_encoder: !new:speechbrain.dataio.encoder.CategoricalEncoder


pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
    loadables:
        mean_var_norm_emb: !ref <mean_var_norm_emb>
        wav2vec2: !ref <wav2vec2>
        model: !ref <model>
        label_encoder: !ref <label_encoder>
    paths:
        mean_var_norm_emb: !ref <pretrained_path>/normalizer_input.ckpt
        wav2vec2: !ref <pretrained_path>/wav2vec2.ckpt
        model: !ref <pretrained_path>/model.ckpt
        label_encoder: !ref <pretrained_path>/label_encoder.txt