nvidia-nemo--speech/examples/speechlm2/salm_eval.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.
import json
from dataclasses import dataclass
from time import perf_counter
from typing import Optional

import lhotse.dataset
import torch
from lhotse import CutSet
from lhotse.serialization import SequentialJsonlWriter
from omegaconf import OmegaConf
from transformers import GenerationConfig
from whisper_normalizer.basic import BasicTextNormalizer
from whisper_normalizer.english import EnglishTextNormalizer

from nemo.collections.asr.metrics.wer import word_error_rate_detail
from nemo.collections.common.data.lhotse.cutset import guess_parse_cutset
from nemo.collections.speechlm2.models import SALM, SALMWithAsrDecoder
from nemo.core.config import hydra_runner
from nemo.utils import logging
from nemo.utils.get_rank import is_global_rank_zero


class ToAudio(torch.utils.data.Dataset):
    def __getitem__(self, cuts: CutSet):
        audios, audio_lens = cuts.load_audio(collate=True)
        return {"cuts": cuts, "audios": audios, "audio_lens": audio_lens}


def _resolve_model_cls(pretrained_name: str, use_asr_decoder: bool, use_nemo_automodel: bool | None):
    """Pick model class. Auto-detects from config.json when use_nemo_automodel is None."""
    if use_asr_decoder:
        return SALMWithAsrDecoder
    if use_nemo_automodel is None:
        # Auto-detect: peek at config.json
        from transformers.utils import cached_file

        config_path = cached_file(
            pretrained_name,
            "config.json",
            _raise_exceptions_for_missing_entries=False,
            _raise_exceptions_for_connection_errors=False,
        )
        if config_path is not None:
            with open(config_path) as f:
                use_nemo_automodel = json.load(f).get("use_nemo_automodel", False)
        else:
            use_nemo_automodel = False
    if use_nemo_automodel:
        from nemo.collections.speechlm2.models import SALMAutomodel

        return SALMAutomodel
    return SALM


@dataclass
class SalmEvalConfig:
    pretrained_name: str
    inputs: str
    batch_size: int = 64
    max_new_tokens: int = 128
    output_manifest: Optional[str] = "generations.jsonl"
    verbose: bool = True
    use_normalizer: Optional[str] = "english"  # "english", "basic", or "none" / "None"
    device: str = "cuda"
    dtype: str = "bfloat16"
    extra_eos_tokens: Optional[list[str]] = None
    system_prompt: Optional[str] = None
    user_prompt: Optional[str] = None
    enable_thinking: Optional[bool] = None
    use_asr_decoder: bool = False  # set this to True if using SALMWithAsrDecoder
    use_nemo_automodel: Optional[bool] = None  # None = auto-detect from config.json
    # Parallelism sizes for distributed inference (launch with torchrun)
    tp_size: int = 1
    ep_size: int = 1
    pp_size: int = 1
    cp_size: int = 1


@hydra_runner(config_name="SalmEvalConfig", schema=SalmEvalConfig)
def main(cfg: SalmEvalConfig):
    logging.info(f'Hydra config:\n{OmegaConf.to_yaml(cfg)}')

    is_distributed = any(s > 1 for s in [cfg.tp_size, cfg.ep_size, cfg.pp_size, cfg.cp_size])
    model_cls = _resolve_model_cls(cfg.pretrained_name, cfg.use_asr_decoder, cfg.use_nemo_automodel)

    if is_distributed and model_cls is SALM:
        raise RuntimeError(
            "Distributed inference requires SALMAutomodel. Set use_nemo_automodel=true or use a checkpoint "
            "exported from SALMAutomodel."
        )

    if is_distributed:
        from nemo.collections.speechlm2.parts.parallel import setup_distributed

        strategy = setup_distributed(
            tp_size=cfg.tp_size, ep_size=cfg.ep_size, pp_size=cfg.pp_size, cp_size=cfg.cp_size
        )
        model = model_cls.from_pretrained(
            cfg.pretrained_name,
            device_mesh=strategy.device_mesh,
            distributed_config=strategy.distributed_config,
            moe_config=strategy.moe_config,
            moe_mesh=strategy.moe_mesh,
            torch_dtype=cfg.dtype,
        )
    else:
        model = model_cls.from_pretrained(cfg.pretrained_name)
        model = model.to(getattr(torch, cfg.dtype)).to(cfg.device)
    model = model.eval()

    cuts = guess_parse_cutset(cfg.inputs).sort_by_duration()
    dloader = torch.utils.data.DataLoader(
        dataset=ToAudio(),
        # rank=0 world_size=1 hardcoded so lhotse doesn't accidentally auto-split batches in model parallel settings
        sampler=lhotse.dataset.DynamicCutSampler(cuts, max_cuts=cfg.batch_size, rank=0, world_size=1),
        num_workers=1,
        batch_size=None,
    )

    normalizer = {"english": EnglishTextNormalizer(), "basic": BasicTextNormalizer()}.get(
        cfg.use_normalizer, lambda x: x
    )

    eos_tokens = [model.text_eos_id]
    if cfg.extra_eos_tokens is not None:
        for t in cfg.extra_eos_tokens:
            tid = model.tokenizer.token_to_id(t)
            assert tid is not None, f"Token '{t}' is not in the model's vocabulary."
            eos_tokens.append(tid)

    # Construct the prompt from ASR data of the form.
    # Optional system prompt goes first.
    prompt = []
    if cfg.system_prompt is not None:
        prompt.append({"role": "system", "content": cfg.system_prompt})
    # If no user prompt is provided, just use the audio placeholder.
    content = model.audio_locator_tag
    # Otherwise:
    # * if user prompt already has audio placeholder, add it as-is,
    # * if not, append audio placeholder at the end of user prompt
    if cfg.user_prompt is not None:
        content = cfg.user_prompt
        if model.audio_locator_tag not in content:
            content = f"{content} {model.audio_locator_tag}"
    prompt.append({"role": "user", "content": content})

    refs = []
    hyps = []
    input_durations = []
    infer_durations = []
    for batch_idx, batch in enumerate(dloader):
        ts = perf_counter()
        answer_ids = model.generate(
            prompts=[prompt] * len(batch["cuts"]),  # identical prompt for each example
            audios=batch["audios"].to(model.device, non_blocking=True),
            audio_lens=batch["audio_lens"].to(model.device, non_blocking=True),
            generation_config=GenerationConfig(
                max_new_tokens=cfg.max_new_tokens,
                bos_token_id=model.text_bos_id,
                eos_token_id=eos_tokens,
                pad_token_id=model.text_pad_id,
            ),
            enable_thinking=cfg.enable_thinking,
        )
        answer_ids = answer_ids.cpu()
        batch_infer_duration = perf_counter() - ts

        batch_duration = sum(c.duration for c in batch["cuts"])
        batch_refs = [normalizer(cut.supervisions[0].text) for cut in batch["cuts"]]
        batch_hyps = [
            normalizer(model.tokenizer.ids_to_text(parse_hyp(ans, eos_tokens)).strip()) for ans in answer_ids
        ]
        if cfg.verbose:
            batch_wer, _, nins, ndel, nsub = word_error_rate_detail(batch_hyps, batch_refs)
            batch_rtfx = batch_duration / batch_infer_duration
            logging.info(
                f"Batch {batch_idx}: WER={batch_wer:.2%} [ins={nins:.2%} del={ndel:.2%} sub={nsub:.2%}] RTFx={batch_rtfx:.1f}"
            )

        refs.extend(batch_refs)
        hyps.extend(batch_hyps)
        input_durations.append(batch_duration)
        infer_durations.append(batch_infer_duration)

    wer, _, nins, ndel, nsub = word_error_rate_detail(hypotheses=hyps, references=refs, use_cer=False)
    rtfx = sum(input_durations) / sum(infer_durations)
    logging.info(f"WER: {wer:.2%} [ins={nins:.2%} del={ndel:.2%} sub={nsub:.2%}]")
    logging.info(f"RTFx: {rtfx:.1f}")

    with _create_output_writer(cfg.output_manifest) as writer:
        for cut, ref, hyp in zip(cuts, refs, hyps):
            writer.write({"id": cut.id, "duration": cut.duration, "text": ref, "pred_text": hyp})


def parse_hyp(answer: torch.Tensor, eos_tokens: list[int]):
    end = torch.isin(answer, torch.tensor(eos_tokens)).nonzero(as_tuple=True)[0]
    if end.numel() == 0:
        return answer
    end = end[0]
    return answer[:end]


class _NullWriter:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return False

    def write(self, data):
        pass


def _create_output_writer(output_manifest: Optional[str]):
    if output_manifest is None or not is_global_rank_zero():
        return _NullWriter()
    return SequentialJsonlWriter(output_manifest)


if __name__ == '__main__':
    main()