# #################################

# Basic inference parameters for speaker-id. We have first a network that

# computes some embeddings. On the top of that, we employ a classifier.

#

# Author:

#  * Mirco Ravanelli 2021

# #################################



# pretrain folders:

pretrained_path: best_model





# Model parameters

n_mels: 23

sample_rate: 16000

n_classes: 28 # In this case, we have 28 speakers

emb_dim: 512 # dimensionality of the embeddings



# Feature extraction

compute_features: !new:speechbrain.lobes.features.Fbank

    n_mels: !ref <n_mels>



# Mean and std normalization of the input features

mean_var_norm: !new:speechbrain.processing.features.InputNormalization

    norm_type: sentence

    std_norm: False



# To design a custom model, either just edit the simple CustomModel

# class that's listed here, or replace this `!new` call with a line

# pointing to a different file you've defined.

embedding_model: !new:custom_model.Xvector

    in_channels: !ref <n_mels>

    activation: !name:torch.nn.LeakyReLU

    tdnn_blocks: 5

    tdnn_channels: [512, 512, 512, 512, 1500]

    tdnn_kernel_sizes: [5, 3, 3, 1, 1]

    tdnn_dilations: [1, 2, 3, 1, 1]

    lin_neurons: !ref <emb_dim>



classifier: !new:custom_model.Classifier

    input_shape: [null, null, !ref <emb_dim>]

    activation: !name:torch.nn.LeakyReLU

    lin_blocks: 1

    lin_neurons: !ref <emb_dim>

    out_neurons: !ref <n_classes>



label_encoder: !new:speechbrain.dataio.encoder.CategoricalEncoder



# Objects in "modules" dict will have their parameters moved to the correct

# device, as well as having train()/eval() called on them by the Brain class.

modules:

    compute_features: !ref <compute_features>

    embedding_model: !ref <embedding_model>

    classifier: !ref <classifier>

    mean_var_norm: !ref <mean_var_norm>



pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer

    loadables:

        embedding_model: !ref <embedding_model>

        classifier: !ref <classifier>

        label_encoder: !ref <label_encoder>

    paths:

        embedding_model: !ref <pretrained_path>/embedding_model.ckpt

        classifier: !ref <pretrained_path>/classifier.ckpt

        label_encoder: !ref <pretrained_path>/label_encoder.txt