# ################################
# Model: Wav2Vec + DNN + CTC + Softmax
# Authors:
# Gaelle Laperriere 2023
# ################################

wav2vec_url: LeBenchmark/wav2vec2-FR-3K-large

# Feature parameters:
sample_rate: 16000
feats_dim: 1024

# Model parameters:
activation: !name:torch.nn.LeakyReLU
dnn_blocks: 3
dnn_neurons: 512

# Decoding parameters:
blank_index: 0

# Outputs:
output_neurons: 67

# ------ Functions and classes

wav2vec2: !new:speechbrain.lobes.models.huggingface_wav2vec.HuggingFaceWav2Vec2
    source: !ref <wav2vec_url>
    output_norm: True
    freeze: Freeze
    save_path: wav2vec2_checkpoint

enc: !new:speechbrain.lobes.models.VanillaNN.VanillaNN
    input_shape: [null, null, !ref <feats_dim>]
    activation: !ref <activation>
    dnn_blocks: !ref <dnn_blocks>
    dnn_neurons: !ref <dnn_neurons>

output_lin: !new:speechbrain.nnet.linear.Linear
    input_size: !ref <dnn_neurons>
    n_neurons: !ref <output_neurons>
    bias: True

log_softmax: !new:speechbrain.nnet.activations.Softmax
    apply_log: True

ctc_cost: !name:speechbrain.nnet.losses.ctc_loss
    blank_index: !ref <blank_index>

modules:
    wav2vec2: !ref <wav2vec2>
    enc: !ref <enc>
    output_lin: !ref <output_lin>

model: !new:torch.nn.ModuleList
    - [!ref <enc>, !ref <output_lin>]

model_wav2vec2: !new:torch.nn.ModuleList
    - [!ref <wav2vec2>]
    
pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
    loadables:
        wav2vec2: !ref <wav2vec2>
        model: !ref <model>