# ################################
# Model: wav2vec2-FR-7K-large + DNN + CTC
# ################################
#

# BPE parameters
token_type: char  # ["unigram", "bpe", "char"]
character_coverage: 1.0

# Model parameters
# activation: !name:torch.nn.LeakyReLU
wav2vec_output_dim: 1024
dnn_neurons: 1024
freeze_wav2vec: false
wav2vec2_hub: LeBenchmark/wav2vec2-FR-7K-large
# Outputs
output_neurons: 76  # BPE size, index(blank/eos/bos) = 0

# Decoding parameters
# Be sure that the bos and eos index match with the BPEs ones
blank_index: 0
bos_index: 1
eos_index: 2

#
# Functions and classes
#

# Encoder network
enc: &id002 !new:speechbrain.nnet.containers.Sequential
  input_shape: [null, null, 1024]
  linear1: !name:speechbrain.nnet.linear.Linear
    n_neurons: 1024
    bias: true
  bn1: !name:speechbrain.nnet.normalization.BatchNorm1d
  activation: !new:torch.nn.LeakyReLU
  drop: !new:torch.nn.Dropout
    p: 0.15
  linear2: !name:speechbrain.nnet.linear.Linear
    n_neurons: 1024
    bias: true
  bn2: !name:speechbrain.nnet.normalization.BatchNorm1d
  activation2: !new:torch.nn.LeakyReLU
  drop2: !new:torch.nn.Dropout
    p: 0.15
  linear3: !name:speechbrain.nnet.linear.Linear
    n_neurons: 1024
    bias: true
  bn3: !name:speechbrain.nnet.normalization.BatchNorm1d
  activation3: !new:torch.nn.LeakyReLU

wav2vec2: &id001 !new:speechbrain.lobes.models.huggingface_wav2vec.HuggingFaceWav2Vec2
  source: !ref <wav2vec2_hub>
  output_norm: true
  freeze: false
  save_path: model.pt

ctc_lin: &id003 !new:speechbrain.nnet.linear.Linear

  input_size: 1024
  n_neurons: 76

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

ctc_cost: !name:speechbrain.nnet.losses.ctc_loss
  blank_index: 0

modules:
  wav2vec2: *id001
  enc: *id002
  ctc_lin: *id003
model: &id004 !new:torch.nn.ModuleList
- [*id002, *id003]
model_opt_class: !name:torch.optim.Adadelta
  lr: 1.0
  rho: 0.95
  eps: 1.e-8

wav2vec_opt_class: !name:torch.optim.Adam
  lr: 0.0001
  
tokenizer: !new:sentencepiece.SentencePieceProcessor

encoder: !new:speechbrain.nnet.containers.LengthsCapableSequential
    wav2vec2: !ref <wav2vec2>
    enc: !ref <enc>
    ctc_lin: !ref <ctc_lin>
    log_softmax: !ref <log_softmax>

decoding_function: !name:speechbrain.decoders.ctc_greedy_decode
    blank_id: !ref <blank_index>
    
decoding_type: greedy

pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
   loadables:
      wav2vec2: !ref <wav2vec2>
      asr: !ref <asr_model>
      tokenizer: !ref <tokenizer>

modules:
  wav2vec2: *id001
  enc: *id002
  ctc_lin: *id003
  encoder: !ref <encoder>