    def __init__(
        self,
        input_size: int = 80,
        encoder_dim: int = 256,
        output_size: int = 256,
        attention_heads: int = 4,
        num_blocks: int = 12,
        reduce_idx: Optional[Union[int, List[int]]] = 5,
        recover_idx: Optional[Union[int, List[int]]] = 11,
        feed_forward_expansion_factor: int = 4,
        dw_stride: bool = False,
        input_dropout_rate: float = 0.1,
        pos_enc_layer_type: str = "rel_pos",
        time_reduction_layer_type: str = "conv1d",
        do_rel_shift: bool = True,
        feed_forward_dropout_rate: float = 0.1,
        attention_dropout_rate: float = 0.1,
        cnn_module_kernel: int = 31,
        cnn_norm_type: str = "batch_norm",
        dropout: float = 0.1,
        causal: bool = False,
        adaptive_scale: bool = True,
        activation_type: str = "swish",
        init_weights: bool = True,
        global_cmvn: torch.nn.Module = None,
        normalize_before: bool = False,
        use_dynamic_chunk: bool = False,
        concat_after: bool = False,
        static_chunk_size: int = 0,
        use_dynamic_left_chunk: bool = False,
    ):
        """Construct SqueezeformerEncoder

        Args:
            input_size to use_dynamic_chunk, see in Transformer BaseEncoder.
            encoder_dim (int): The hidden dimension of encoder layer.
            output_size (int): The output dimension of final projection layer.
            attention_heads (int): Num of attention head in attention module.
            num_blocks (int): Num of encoder layers.
            reduce_idx Optional[Union[int, List[int]]]:
                reduce layer index, from 40ms to 80ms per frame.
            recover_idx Optional[Union[int, List[int]]]:
                recover layer index, from 80ms to 40ms per frame.
            feed_forward_expansion_factor (int): Enlarge coefficient of FFN.
            dw_stride (bool): Whether do depthwise convolution
                              on subsampling module.
            input_dropout_rate (float): Dropout rate of input projection layer.
            pos_enc_layer_type (str): Self attention type.
            time_reduction_layer_type (str): Conv1d or Conv2d reduction layer.
            do_rel_shift (bool): Whether to do relative shift
                                 operation on rel-attention module.
            cnn_module_kernel (int): Kernel size of CNN module.
            activation_type (str): Encoder activation function type.
            use_cnn_module (bool): Whether to use convolution module.
            cnn_module_kernel (int): Kernel size of convolution module.
            adaptive_scale (bool): Whether to use adaptive scale.
            init_weights (bool): Whether to initialize weights.
            causal (bool): whether to use causal convolution or not.
        """
        super(SqueezeformerEncoder, self).__init__()
        self.global_cmvn = global_cmvn
        self.reduce_idx: Optional[Union[int, List[int]]] = (
            [reduce_idx] if type(reduce_idx) == int else reduce_idx
        )
        self.recover_idx: Optional[Union[int, List[int]]] = (
            [recover_idx] if type(recover_idx) == int else recover_idx
        )
        self.check_ascending_list()
        if reduce_idx is None:
            self.time_reduce = None
        else:
            if recover_idx is None:
                self.time_reduce = "normal"  # no recovery at the end
            else:
                self.time_reduce = "recover"  # recovery at the end
                assert len(self.reduce_idx) == len(self.recover_idx)
            self.reduce_stride = 2
        self._output_size = output_size
        self.normalize_before = normalize_before
        self.static_chunk_size = static_chunk_size
        self.use_dynamic_chunk = use_dynamic_chunk
        self.use_dynamic_left_chunk = use_dynamic_left_chunk
        self.pos_enc_layer_type = pos_enc_layer_type
        activation = get_activation(activation_type)

        # self-attention module definition
        if pos_enc_layer_type != "rel_pos":
            encoder_selfattn_layer = MultiHeadedAttention
            encoder_selfattn_layer_args = (
                attention_heads,
                output_size,
                attention_dropout_rate,
            )
        else:
            encoder_selfattn_layer = RelPositionMultiHeadedAttention
            encoder_selfattn_layer_args = (
                attention_heads,
                encoder_dim,
                attention_dropout_rate,
                do_rel_shift,
                adaptive_scale,
                init_weights,
            )

        # feed-forward module definition
        positionwise_layer = PositionwiseFeedForward
        positionwise_layer_args = (
            encoder_dim,
            encoder_dim * feed_forward_expansion_factor,
            feed_forward_dropout_rate,
            activation,
            adaptive_scale,
            init_weights,
        )

        # convolution module definition
        convolution_layer = ConvolutionModule
        convolution_layer_args = (
            encoder_dim,
            cnn_module_kernel,
            activation,
            cnn_norm_type,
            causal,
            True,
            adaptive_scale,
            init_weights,
        )

        self.embed = DepthwiseConv2dSubsampling4(
            1,
            encoder_dim,
            RelPositionalEncoding(encoder_dim, dropout_rate=0.1),
            dw_stride,
            input_size,
            input_dropout_rate,
            init_weights,
        )

        self.preln = nn.LayerNorm(encoder_dim)
        self.encoders = torch.nn.ModuleList(
            [
                SqueezeformerEncoderLayer(
                    encoder_dim,
                    encoder_selfattn_layer(*encoder_selfattn_layer_args),
                    positionwise_layer(*positionwise_layer_args),
                    convolution_layer(*convolution_layer_args),
                    positionwise_layer(*positionwise_layer_args),
                    normalize_before,
                    dropout,
                    concat_after,
                )
                for _ in range(num_blocks)
            ]
        )
        if time_reduction_layer_type == "conv1d":
            time_reduction_layer = TimeReductionLayer1D
            time_reduction_layer_args = {
                "channel": encoder_dim,
                "out_dim": encoder_dim,
            }
        elif time_reduction_layer_type == "stream":
            time_reduction_layer = TimeReductionLayerStream
            time_reduction_layer_args = {
                "channel": encoder_dim,
                "out_dim": encoder_dim,
            }
        else:
            time_reduction_layer = TimeReductionLayer2D
            time_reduction_layer_args = {"encoder_dim": encoder_dim}

        self.time_reduction_layer = time_reduction_layer(**time_reduction_layer_args)
        self.time_recover_layer = nn.Linear(encoder_dim, encoder_dim)
        self.final_proj = None
        if output_size != encoder_dim:
            self.final_proj = nn.Linear(encoder_dim, output_size)
Construct SqueezeformerEncoder

Args:
input_size to use_dynamic_chunk, see in Transformer BaseEncoder.
encoder_dim (int): The hidden dimension of encoder layer.
output_size (int): The output dimension of final projection layer.
attention_heads (int): Num of attention head in attention module.
num_blocks (int): Num of encoder layers.
reduce_idx Optional[Union[int, List[int]]]:
reduce layer index, from 40ms to 80ms per frame.
recover_idx Optional[Union[int, List[int]]]:
recover layer index, from 80ms to 40ms per frame.
feed_forward_expansion_factor (int): Enlarge coefficient of FFN.
dw_stride (bool): Whether do depthwise convolution
on subsampling module.
input_dropout_rate (float): Dropout rate of input projection layer.
pos_enc_layer_type (str): Self attention type.
time_reduction_layer_type (str): Conv1d or Conv2d reduction layer.
do_rel_shift (bool): Whether to do relative shift
operation on rel-attention module.
cnn_module_kernel (int): Kernel size of CNN module.
activation_type (str): Encoder activation function type.
use_cnn_module (bool): Whether to use convolution module.
cnn_module_kernel (int): Kernel size of convolution module.
adaptive_scale (bool): Whether to use adaptive scale.
init_weights (bool): Whether to initialize weights.
causal (bool): whether to use causal convolution or not.
modules/wenet_extractor/squeezeformer/encoder.py