# ############################################################################
# Model: RE-SepFormer for source separation
# Dataset: Libri2mix (16k Sample Rate)
# ############################################################################

# Basic parameters
sample_rate: 16000
num_spks: 2

# Encoder parameters
N_encoder_out: 128
out_channels: 128
kernel_size: 16
kernel_stride: 8
d_ffn: 1024

# Specifying the network
Encoder: !new:speechbrain.lobes.models.dual_path.Encoder &id004
    kernel_size: 16
    out_channels: 128

# RE-SepFormer specific Transformer Blocks
intra_mdl:
    !new:speechbrain.lobes.models.resepformer.SBTransformerBlock_wnormandskip &id002
    num_layers: 8
    d_model: 128
    nhead: 8
    d_ffn: 1024
    dropout: 0
    use_positional_encoding: true
    norm_before: true
    use_norm: true
    use_skip: true

mem_mdl:
    !new:speechbrain.lobes.models.resepformer.SBTransformerBlock_wnormandskip &id003
    num_layers: 8
    d_model: 128
    nhead: 8
    d_ffn: 1024
    dropout: 0
    use_positional_encoding: true
    norm_before: true
    use_norm: true
    use_skip: true

MaskNet:
    !new:speechbrain.lobes.models.resepformer.ResourceEfficientSeparator &id006
    input_dim: 128
    num_spk: 2
    causal: false
    unit: 256
    segment_size: 150
    layer: 2
    mem_type: av
    seg_model: *id002
    mem_model: *id003

Decoder: !new:speechbrain.lobes.models.dual_path.Decoder &id005
    in_channels: 128
    out_channels: 1
    kernel_size: 16
    stride: 8
    bias: false

modules:
    encoder: *id004
    decoder: *id005
    masknet: *id006

pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
    loadables:
        encoder: !ref <Encoder>
        decoder: !ref <Decoder>
        masknet: !ref <MaskNet>