FunASR/funasr/models/encoder/chunk_encoder_utils/building.py

"""Set of methods to build Transducer encoder architecture."""

from typing import Any, Dict, List, Optional, Union

from funasr.modules.activation import get_activation
from funasr.models.encoder.chunk_encoder_blocks.branchformer import Branchformer
from funasr.models.encoder.chunk_encoder_blocks.conformer import Conformer
from funasr.models.encoder.chunk_encoder_blocks.conv1d import Conv1d
from funasr.models.encoder.chunk_encoder_blocks.conv_input import ConvInput
from funasr.models.encoder.chunk_encoder_blocks.linear_input import LinearInput
from funasr.models.encoder.chunk_encoder_modules.attention import (  # noqa: H301
    RelPositionMultiHeadedAttention,
)
from funasr.models.encoder.chunk_encoder_modules.convolution import (  # noqa: H301
    ConformerConvolution,
    ConvolutionalSpatialGatingUnit,
)
from funasr.models.encoder.chunk_encoder_modules.multi_blocks import MultiBlocks
from funasr.models.encoder.chunk_encoder_modules.normalization import get_normalization
from funasr.models.encoder.chunk_encoder_modules.positional_encoding import (  # noqa: H301
    RelPositionalEncoding,
)
from funasr.modules.positionwise_feed_forward import (
    PositionwiseFeedForward,
)


def build_main_parameters(
    pos_wise_act_type: str = "swish",
    conv_mod_act_type: str = "swish",
    pos_enc_dropout_rate: float = 0.0,
    pos_enc_max_len: int = 5000,
    simplified_att_score: bool = False,
    norm_type: str = "layer_norm",
    conv_mod_norm_type: str = "layer_norm",
    after_norm_eps: Optional[float] = None,
    after_norm_partial: Optional[float] = None,
    dynamic_chunk_training: bool = False,
    short_chunk_threshold: float = 0.75,
    short_chunk_size: int = 25,
    left_chunk_size: int = 0,
    time_reduction_factor: int = 1,
    unified_model_training: bool = False,
    default_chunk_size: int = 16,
    jitter_range: int =4,
    **activation_parameters,
) -> Dict[str, Any]:
    """Build encoder main parameters.

    Args:
        pos_wise_act_type: Conformer position-wise feed-forward activation type.
        conv_mod_act_type: Conformer convolution module activation type.
        pos_enc_dropout_rate: Positional encoding dropout rate.
        pos_enc_max_len: Positional encoding maximum length.
        simplified_att_score: Whether to use simplified attention score computation.
        norm_type: X-former normalization module type.
        conv_mod_norm_type: Conformer convolution module normalization type.
        after_norm_eps: Epsilon value for the final normalization.
        after_norm_partial: Value for the final normalization with RMSNorm.
        dynamic_chunk_training: Whether to use dynamic chunk training.
        short_chunk_threshold: Threshold for dynamic chunk selection.
        short_chunk_size: Minimum number of frames during dynamic chunk training.
        left_chunk_size: Number of frames in left context.
        **activations_parameters: Parameters of the activation functions.
                                    (See espnet2/asr_transducer/activation.py)

    Returns:
        : Main encoder parameters

    """
    main_params = {}

    main_params["pos_wise_act"] = get_activation(
        pos_wise_act_type, **activation_parameters
    )

    main_params["conv_mod_act"] = get_activation(
        conv_mod_act_type, **activation_parameters
    )

    main_params["pos_enc_dropout_rate"] = pos_enc_dropout_rate
    main_params["pos_enc_max_len"] = pos_enc_max_len

    main_params["simplified_att_score"] = simplified_att_score

    main_params["norm_type"] = norm_type
    main_params["conv_mod_norm_type"] = conv_mod_norm_type

    (
        main_params["after_norm_class"],
        main_params["after_norm_args"],
    ) = get_normalization(norm_type, eps=after_norm_eps, partial=after_norm_partial)

    main_params["dynamic_chunk_training"] = dynamic_chunk_training
    main_params["short_chunk_threshold"] = max(0, short_chunk_threshold)
    main_params["short_chunk_size"] = max(0, short_chunk_size)
    main_params["left_chunk_size"] = max(0, left_chunk_size)
    
    main_params["unified_model_training"] = unified_model_training
    main_params["default_chunk_size"] = max(0, default_chunk_size)
    main_params["jitter_range"] = max(0, jitter_range)
   
    main_params["time_reduction_factor"] = time_reduction_factor

    return main_params


def build_positional_encoding(
    block_size: int, configuration: Dict[str, Any]
) -> RelPositionalEncoding:
    """Build positional encoding block.

    Args:
        block_size: Input/output size.
        configuration: Positional encoding configuration.

    Returns:
        : Positional encoding module.

    """
    return RelPositionalEncoding(
        block_size,
        configuration.get("pos_enc_dropout_rate", 0.0),
        max_len=configuration.get("pos_enc_max_len", 5000),
    )


def build_input_block(
    input_size: int,
    configuration: Dict[str, Union[str, int]],
) -> ConvInput:
    """Build encoder input block.

    Args:
        input_size: Input size.
        configuration: Input block configuration.

    Returns:
        : ConvInput block function.

    """
    if configuration["linear"]:
        return LinearInput(
            input_size,
            configuration["output_size"],
            configuration["subsampling_factor"],
        )
    else:
        return ConvInput(
            input_size,
            configuration["conv_size"],
            configuration["subsampling_factor"],
            vgg_like=configuration["vgg_like"],
            output_size=configuration["output_size"],
        )


def build_branchformer_block(
    configuration: List[Dict[str, Any]],
    main_params: Dict[str, Any],
) -> Conformer:
    """Build Branchformer block.

    Args:
        configuration: Branchformer block configuration.
        main_params: Encoder main parameters.

    Returns:
        : Branchformer block function.

    """
    hidden_size = configuration["hidden_size"]
    linear_size = configuration["linear_size"]

    dropout_rate = configuration.get("dropout_rate", 0.0)

    conv_mod_norm_class, conv_mod_norm_args = get_normalization(
        main_params["conv_mod_norm_type"],
        eps=configuration.get("conv_mod_norm_eps"),
        partial=configuration.get("conv_mod_norm_partial"),
    )

    conv_mod_args = (
        linear_size,
        configuration["conv_mod_kernel_size"],
        conv_mod_norm_class,
        conv_mod_norm_args,
        dropout_rate,
        main_params["dynamic_chunk_training"],
    )

    mult_att_args = (
        configuration.get("heads", 4),
        hidden_size,
        configuration.get("att_dropout_rate", 0.0),
        main_params["simplified_att_score"],
    )

    norm_class, norm_args = get_normalization(
        main_params["norm_type"],
        eps=configuration.get("norm_eps"),
        partial=configuration.get("norm_partial"),
    )

    return lambda: Branchformer(
        hidden_size,
        linear_size,
        RelPositionMultiHeadedAttention(*mult_att_args),
        ConvolutionalSpatialGatingUnit(*conv_mod_args),
        norm_class=norm_class,
        norm_args=norm_args,
        dropout_rate=dropout_rate,
    )


def build_conformer_block(
    configuration: List[Dict[str, Any]],
    main_params: Dict[str, Any],
) -> Conformer:
    """Build Conformer block.

    Args:
        configuration: Conformer block configuration.
        main_params: Encoder main parameters.

    Returns:
        : Conformer block function.

    """
    hidden_size = configuration["hidden_size"]
    linear_size = configuration["linear_size"]

    pos_wise_args = (
        hidden_size,
        linear_size,
        configuration.get("pos_wise_dropout_rate", 0.0),
        main_params["pos_wise_act"],
    )

    conv_mod_norm_args = {
        "eps": configuration.get("conv_mod_norm_eps", 1e-05),
        "momentum": configuration.get("conv_mod_norm_momentum", 0.1),
    }

    conv_mod_args = (
        hidden_size,
        configuration["conv_mod_kernel_size"],
        main_params["conv_mod_act"],
        conv_mod_norm_args,
        main_params["dynamic_chunk_training"] or main_params["unified_model_training"],
    )

    mult_att_args = (
        configuration.get("heads", 4),
        hidden_size,
        configuration.get("att_dropout_rate", 0.0),
        main_params["simplified_att_score"],
    )

    norm_class, norm_args = get_normalization(
        main_params["norm_type"],
        eps=configuration.get("norm_eps"),
        partial=configuration.get("norm_partial"),
    )

    return lambda: Conformer(
        hidden_size,
        RelPositionMultiHeadedAttention(*mult_att_args),
        PositionwiseFeedForward(*pos_wise_args),
        PositionwiseFeedForward(*pos_wise_args),
        ConformerConvolution(*conv_mod_args),
        norm_class=norm_class,
        norm_args=norm_args,
        dropout_rate=configuration.get("dropout_rate", 0.0),
    )


def build_conv1d_block(
    configuration: List[Dict[str, Any]],
    causal: bool,
) -> Conv1d:
    """Build Conv1d block.

    Args:
        configuration: Conv1d block configuration.

    Returns:
        : Conv1d block function.

    """
    return lambda: Conv1d(
        configuration["input_size"],
        configuration["output_size"],
        configuration["kernel_size"],
        stride=configuration.get("stride", 1),
        dilation=configuration.get("dilation", 1),
        groups=configuration.get("groups", 1),
        bias=configuration.get("bias", True),
        relu=configuration.get("relu", True),
        batch_norm=configuration.get("batch_norm", False),
        causal=causal,
        dropout_rate=configuration.get("dropout_rate", 0.0),
    )


def build_body_blocks(
    configuration: List[Dict[str, Any]],
    main_params: Dict[str, Any],
    output_size: int,
) -> MultiBlocks:
    """Build encoder body blocks.

    Args:
        configuration: Body blocks configuration.
        main_params: Encoder main parameters.
        output_size: Architecture output size.

    Returns:
        MultiBlocks function encapsulation all encoder blocks.

    """
    fn_modules = []
    extended_conf = []

    for c in configuration:
        if c.get("num_blocks") is not None:
            extended_conf += c["num_blocks"] * [
                {c_i: c[c_i] for c_i in c if c_i != "num_blocks"}
            ]
        else:
            extended_conf += [c]

    for i, c in enumerate(extended_conf):
        block_type = c["block_type"]

        if block_type == "branchformer":
            module = build_branchformer_block(c, main_params)
        elif block_type == "conformer":
            module = build_conformer_block(c, main_params)
        elif block_type == "conv1d":
            module = build_conv1d_block(c, main_params["dynamic_chunk_training"])
        else:
            raise NotImplementedError

        fn_modules.append(module)

    return MultiBlocks(
        [fn() for fn in fn_modules],
        output_size,
        norm_class=main_params["after_norm_class"],
        norm_args=main_params["after_norm_args"],
    )