sample_rate: 16000

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

# Model parameters
activation: !name:torch.nn.GELU
wav2vec_output_dim: 1024
dnn_neurons: 1024
freeze_wav2vec: false
dropout: 0.2

# Outputs
output_neurons: 80 # BPE size, index(blank/eos/bos) = 0
tokenizer: !new:sentencepiece.SentencePieceProcessor

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

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.GELU
  drop: !new:torch.nn.Dropout
    p: 0.2
  linear2: !name:speechbrain.nnet.linear.Linear
    n_neurons: 1024
    bias: true
  bn2: !name:speechbrain.nnet.normalization.BatchNorm1d
  activation2: !new:torch.nn.GELU
  drop2: !new:torch.nn.Dropout
    p: 0.2
  linear3: !name:speechbrain.nnet.linear.Linear
    n_neurons: 1024
    bias: true
  bn3: !name:speechbrain.nnet.normalization.BatchNorm1d
  activation3: !new:torch.nn.GELU

wav2vec2: !new:speechbrain.lobes.models.huggingface_wav2vec.HuggingFaceWav2Vec2
  source: asafaya/hubert-large-arabic-transcribe
  output_norm: true
  freeze: false
  save_path: wav2vec2_checkpoint

ctc_lin: !new:speechbrain.nnet.linear.Linear
  input_size: 1024
  n_neurons: 80

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

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

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

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

error_rate_computer: !name:speechbrain.utils.metric_stats.ErrorRateStats
cer_computer: !name:speechbrain.utils.metric_stats.ErrorRateStats
  split_tokens: true

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

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