# ################################
# Model: Fastspeech2 for TTS
# Authors: Sathvik Udupa, Yingzhi Wang 
# ################################
n_symbols: 62 #fixed deppending on symbols in textToSequence
n_mel_channels: 80
padding_idx: 0

# Encoder parameters
enc_num_layers: 4
enc_num_head: 2
enc_d_model: 384
enc_ffn_dim: 1024
enc_k_dim: 384
enc_v_dim: 384
enc_dropout: 0.1

# Decoder parameters
dec_num_layers: 4
dec_num_head: 2
dec_d_model: 384
dec_ffn_dim: 1024
dec_k_dim: 384
dec_v_dim: 384
dec_dropout: 0.1

# common
normalize_before: True
ffn_type: 1dcnn #1dcnn or ffn
dur_pred_kernel_size: 3
pitch_pred_kernel_size: 3
energy_pred_kernel_size: 3

model: !new:speechbrain.lobes.models.FastSpeech2.FastSpeech2
    enc_num_layers: !ref <enc_num_layers>
    enc_num_head: !ref <enc_num_head>
    enc_d_model: !ref <enc_d_model>
    enc_ffn_dim: !ref <enc_ffn_dim>
    enc_k_dim: !ref <enc_k_dim>
    enc_v_dim: !ref <enc_v_dim>
    enc_dropout: !ref <enc_dropout>
    dec_num_layers: !ref <dec_num_layers>
    dec_num_head: !ref <dec_num_head>
    dec_d_model: !ref <dec_d_model>
    dec_ffn_dim: !ref <dec_ffn_dim>
    dec_k_dim: !ref <dec_k_dim>
    dec_v_dim: !ref <dec_v_dim>
    dec_dropout: !ref <dec_dropout>
    normalize_before: !ref <normalize_before>
    ffn_type: !ref <ffn_type>
    n_char: !ref <n_symbols>
    n_mels: !ref <n_mel_channels>
    padding_idx: !ref <padding_idx>
    dur_pred_kernel_size: !ref <dur_pred_kernel_size>
    pitch_pred_kernel_size: !ref <pitch_pred_kernel_size>
    energy_pred_kernel_size: !ref <energy_pred_kernel_size>

# The lexicon file must be the same used for training
lexicon:
    - "t"
    - "?"
    - "q"
    - "j"
    - "g"
    - "p"
    - "x"
    - "("
    - "é"
    - "e"
    - "z"
    - ","
    - "o"
    - "a"
    - "m"
    - "n"
    - "u"
    - "d"
    - ":"
    - "w"
    - "à"
    - "“"
    - "."
    - "”"
    - "’"
    - "["
    - "v"
    - "h"
    - " " 
    - "ê"
    - "b"
    - "'"
    - "\""
    - "f"
    - "â"
    - "!"
    - ";"
    - "l"
    - "r"
    - "è"
    - "i"
    - "]"
    - "s"
    - "k"
    - "y"
    - ")"
    - "c"
    - "ü"
    - "-"


input_encoder: !new:speechbrain.dataio.encoder.TextEncoder


modules:
    model: !ref <model>

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