# ############################################################################
# Model: WAV2VEC base for Emotion Recognition
# ############################################################################


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

# Feature parameters
wav2vec2_hub: facebook\wav2vecChinese

# Pretrain folder (HuggingFace)
pretrained_path: emotion-recognition-wav2vec2-IEMOCAP

# parameters
encoder_dim: 1024
out_n_neurons: 2

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

LSTM: !new:speechbrain.nnet.RNN.LSTM
    input_size: !ref <encoder_dim>
    hidden_size: !ref <encoder_dim>

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:
    wav2vec2: !ref <wav2vec2>
    output_mlp: !ref <output_mlp>
    avg_pool: !ref <avg_pool>
    LSTM: !ref <LSTM>

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


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