# ############################################################################
# Model: SENSE multilingual semantic speech encoder
# ############################################################################

# Audio configuration
sample_rate: 16000

# Hugging Face repository containing the pretrained SENSE checkpoints
pretrained_path: LIA-AvignonUniversity/SENSE


# ---------------------------------------------------------------------------
# Audio normalization
# ---------------------------------------------------------------------------

audio_normalizer: !new:speechbrain.dataio.preprocess.AudioNormalizer
    sample_rate: !ref <sample_rate>


# ---------------------------------------------------------------------------
# Speech encoder: w2v-BERT 2.0
# ---------------------------------------------------------------------------

wav2vec2_hub: facebook/w2v-bert-2.0

wav2vec2: !new:speechbrain.integrations.huggingface.w2v_bert.W2VBert
    source: !ref <wav2vec2_hub>
    save_path: wav2vec2_checkpoints
    freeze: True
    sample_rate: !ref <sample_rate>


# ---------------------------------------------------------------------------
# Utterance-level attention pooling
# ---------------------------------------------------------------------------

d_model: 1024

attn_pooling: !new:speechbrain.nnet.pooling.AttentionPooling
    input_dim: !ref <d_model>


# ---------------------------------------------------------------------------
# Attention pooling module stored in model.ckpt
# ---------------------------------------------------------------------------

model: !new:torch.nn.ModuleList
    modules:
        - !ref <attn_pooling>


# ---------------------------------------------------------------------------
# Modules exposed to SENSEEncoder
# ---------------------------------------------------------------------------

modules:
    wav2vec2: !ref <wav2vec2>
    attn_pooling: !ref <attn_pooling>


# ---------------------------------------------------------------------------
# Pretrained checkpoint loading
# ---------------------------------------------------------------------------

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

    paths:
        wav2vec2: !ref <pretrained_path>/wav2vec2.ckpt
        model: !ref <pretrained_path>/model.ckpt