/nemo/collections/asr/inference/factory/base_builder.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from omegaconf import OmegaConf, open_dict
from omegaconf.dictconfig import DictConfig

from nemo.collections.asr.inference.model_wrappers.cache_aware_ctc_inference_wrapper import (
    CacheAwareCTCInferenceWrapper,
)
from nemo.collections.asr.inference.model_wrappers.cache_aware_rnnt_inference_wrapper import (
    CacheAwareRNNTInferenceWrapper,
)
from nemo.collections.asr.inference.model_wrappers.ctc_inference_wrapper import CTCInferenceWrapper
from nemo.collections.asr.inference.model_wrappers.rnnt_inference_wrapper import RNNTInferenceWrapper
from nemo.collections.asr.inference.model_wrappers.salm_asr_inference_wrapper import SALMASRInferenceWrapper
from nemo.collections.asr.inference.utils.enums import ASRDecodingType, PipelineType
from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
from nemo.utils import logging

if TYPE_CHECKING:
    from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
    from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator


class BaseBuilder:
    """
    Base Builder class.
    Builds the ASR/ITN components.
    Derived classes should implement the `build` method which should include the logic of creating concrete pipeline.
    """

    @classmethod
    def _build_nmt(cls, cfg: DictConfig) -> LLMTranslator | None:
        """
        Build the NMT model based on the config.
        Args:
            cfg: (DictConfig) Config
        Returns:
            (LLMTranslator | None) NMT model
        """
        nmt_model = None
        if cfg.enable_nmt:
            from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator

            nmt_model = LLMTranslator(
                model_name=cfg.nmt.model_name,
                source_language=cfg.nmt.source_language,
                target_language=cfg.nmt.target_language,
                waitk=cfg.nmt.waitk,
                device=cfg.nmt.device,
                device_id=cfg.nmt.device_id,
                batch_size=cfg.nmt.batch_size,
                llm_params=cfg.nmt.llm_params,
                sampling_params=cfg.nmt.sampling_params,
            )
            logging.info(f"NMT model `{cfg.nmt.model_name}` loaded")
        return nmt_model

    @staticmethod
    def _apply_confidence_cfg(cfg: DictConfig, decoding_cfg: RNNTDecodingConfig) -> None:
        """
        Wire the separately-stored `confidence` block into the RNNT decoding confidence config so the
        greedy or batched-beam decoder computes per-token confidence with the configured method. The streaming
        pipelines only support non-blank confidence (`confidence.exclude_blank=true`).
        Args:
            cfg: (DictConfig) Full pipeline config (provides the top-level `confidence` block).
            decoding_cfg: (RNNTDecodingConfig) Decoding config to update in place.
        """
        preserve_frame_confidence = decoding_cfg.greedy.get(
            "preserve_frame_confidence", False
        ) or decoding_cfg.beam.get("preserve_frame_confidence", False)
        if not preserve_frame_confidence:
            return
        confidence_cfg = cfg.get("confidence", None)
        if confidence_cfg is None:
            return
        if not confidence_cfg.get("exclude_blank", True):
            raise ValueError(
                "Streaming confidence supports only non-blank confidence (`confidence.exclude_blank=true`)."
            )
        decoding_cfg.confidence_cfg.preserve_frame_confidence = True
        decoding_cfg.confidence_cfg.preserve_token_confidence = True
        decoding_cfg.confidence_cfg.preserve_word_confidence = True
        decoding_cfg.confidence_cfg.exclude_blank = True
        decoding_cfg.confidence_cfg.aggregation = confidence_cfg.get("aggregation", "mean")
        decoding_cfg.confidence_cfg.method_cfg = OmegaConf.merge(
            decoding_cfg.confidence_cfg.method_cfg, confidence_cfg.method_cfg
        )

    @classmethod
    def _build_asr(cls, cfg: DictConfig, decoding_cfg: CTCDecodingConfig | RNNTDecodingConfig | None) -> Any:
        """
        Build the ASR model based on the config.
        Args:
            cfg: (DictConfig) Config
            decoding_cfg: (CTCDecodingConfig | RNNTDecodingConfig | None) Decoding config
        Returns:
            (Any) ASR inference model
        """

        asr_decoding_type = ASRDecodingType.from_str(cfg.asr_decoding_type)
        pipeline_type = PipelineType.from_str(cfg.pipeline_type)
        model_params = {
            "model_name": cfg.asr.model_name,
            "device": cfg.asr.device,
            "device_id": cfg.asr.device_id,
            "compute_dtype": cfg.asr.compute_dtype,
            "use_amp": cfg.asr.use_amp,
            "decoding_cfg": decoding_cfg,
        }
        match (asr_decoding_type, pipeline_type):
            case (ASRDecodingType.CTC, PipelineType.BUFFERED):
                asr_class = CTCInferenceWrapper
            case (ASRDecodingType.RNNT, PipelineType.BUFFERED):
                asr_class = RNNTInferenceWrapper
            case (ASRDecodingType.SALM, PipelineType.BUFFERED):
                asr_class = SALMASRInferenceWrapper
                # remove decoding_cfg, SALM AED does not use decoding_cfg yet
                model_params.pop("decoding_cfg")
            case (ASRDecodingType.CTC, PipelineType.CACHE_AWARE):
                asr_class = CacheAwareCTCInferenceWrapper
            case (ASRDecodingType.RNNT, PipelineType.CACHE_AWARE):
                asr_class = CacheAwareRNNTInferenceWrapper
            case _:
                raise ValueError(
                    f"Wrong combination of ASR decoding type and pipeline type: {asr_decoding_type, pipeline_type}"
                )

        asr_model = asr_class(**model_params)
        logging.info(f"ASR model `{cfg.asr.model_name}` loaded")
        return asr_model

    @classmethod
    def _build_itn(cls, cfg: DictConfig, input_is_lower_cased: bool) -> AlignmentPreservingInverseNormalizer | None:
        """
        Build the ITN model based on the config.
        Args:
            cfg: (DictConfig) Config
            input_is_lower_cased: (bool) Whether the input is lower cased
        Returns:
            (AlignmentPreservingInverseNormalizer | None) ITN model
        """
        itn_model = None
        if cfg.enable_itn:
            # Do not remove this import. It is used to avoid nemo_text_processing import when verbatim transcripts is enabled.
            from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer

            input_case = (
                AlignmentPreservingInverseNormalizer.LOWER_CASED
                if input_is_lower_cased
                else AlignmentPreservingInverseNormalizer.UPPER_CASED
            )

            target_lang = getattr(cfg, "lang", getattr(cfg, "target_lang", None))
            if target_lang is None:
                raise ValueError("Language is not specified. Cannot load ITN model.")

            itn_cfg = cfg.itn
            with open_dict(itn_cfg):
                itn_cfg.lang = target_lang
                itn_cfg.input_case = input_case
                itn_cfg.cache_dir = cfg.cache_dir

            itn_model = AlignmentPreservingInverseNormalizer(
                lang=itn_cfg.lang,
                input_case=itn_cfg.input_case,
                whitelist=itn_cfg.whitelist,
                cache_dir=itn_cfg.cache_dir,
                overwrite_cache=itn_cfg.overwrite_cache,
                max_number_of_permutations_per_split=itn_cfg.max_number_of_permutations_per_split,
            )
            logging.info(f"Built inverse text normalizer with the input case: `{input_case}`.")

        if itn_model is not None:
            logging.info("ITN model loaded")
        return itn_model

    @classmethod
    def build(cls, cfg: DictConfig) -> Any:
        """
        Build the pipeline based on the config.
        Args:
            cfg: (DictConfig) Config
        Returns:
            Returns object responsible for the inference
        """
        raise NotImplementedError("This method should be implemented in subclasses.")