diff --git a/examples/asr/asr_streaming_inference/run_nemo_simulstream.py b/examples/asr/asr_streaming_inference/run_nemo_simulstream.py new file mode 100644 index 000000000000..0794956d5066 --- /dev/null +++ b/examples/asr/asr_streaming_inference/run_nemo_simulstream.py @@ -0,0 +1,281 @@ +# Copyright (c) 2026, 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. + +""" +Wrapper to run simulstream (https://github.com/NVIDIA/simulstream) inference with NeMo's streaming +ASR/AST configs (see nemo/collections/asr/inference/utils/simulstream_pipeline_adapter.py). + +NeMo config files don't have the 'type' field simulstream requires to locate the speech processor +class; this script adds it (and a couple of other simulstream-required fields) on-the-fly. + +Requires `pip install simulstream` (not a NeMo dependency). + +Usage with a wav list: + python run_nemo_simulstream.py \\ + --config path/to/cache_aware_rnnt.yaml \\ + --wav-list audio_list.txt \\ + --src-lang ru \\ + --tgt-lang en \\ + --metrics-log metrics.jsonl + +Usage with a NeMo manifest: + python run_nemo_simulstream.py \\ + --config path/to/cache_aware_rnnt.yaml \\ + --manifest data/manifest.json \\ + --src-lang ru \\ + --tgt-lang en \\ + --metrics-log metrics.jsonl + +Any additional `key=value` arguments are applied as NeMo config overrides (e.g. `asr.device_id=1`). +""" + +import argparse +import shutil +import subprocess +import tempfile +from pathlib import Path + +from omegaconf import OmegaConf + +from nemo.collections.asr.inference.utils.simulstream_manifest_utils import load_manifest_audio_paths +from nemo.utils import logging + +# ISO 639-1 code -> full name for NeMo's NMT config. Extend as needed. +LANGUAGE_CODES = { + "bg": "Bulgarian", + "hr": "Croatian", + "cs": "Czech", + "da": "Danish", + "nl": "Dutch", + "en": "English", + "et": "Estonian", + "fi": "Finnish", + "fr": "French", + "de": "German", + "el": "Greek", + "hu": "Hungarian", + "it": "Italian", + "lv": "Latvian", + "lt": "Lithuanian", + "mt": "Maltese", + "pl": "Polish", + "pt": "Portuguese", + "ro": "Romanian", + "sk": "Slovak", + "sl": "Slovenian", + "es": "Spanish", + "sv": "Swedish", + "ru": "Russian", + "uk": "Ukrainian", +} + +# Languages that don't separate words with whitespace, so simulstream latency (and any other +# word-count-based metric) must be computed at the character level instead. Extend as needed. +LATENCY_UNIT_CHAR_LANGUAGES = {"zh", "ja", "ko", "th", "lo", "my", "km"} + + +def get_language_name(code: str) -> str: + """Map a language code to its full name for the NeMo NMT config; unknown codes pass through.""" + return LANGUAGE_CODES.get(code, code) + + +def get_latency_unit(code: str) -> str: + """Map a language code to its latency unit for simulstream metrics; unknown codes default to 'word'.""" + return "char" if code in LATENCY_UNIT_CHAR_LANGUAGES else "word" + + +def add_simulstream_fields( + cfg_path: str, + metrics_log: str, + src_lang: str = None, + tgt_lang: str = None, + overrides: list = None, + reference_manifest: str | None = None, + output_manifest: str | None = None, +) -> str: + """ + Load a NeMo config and add the fields simulstream requires. + + The simulstream `speech_chunk_size` is always derived from the NeMo config + (`streaming.chunk_size` for buffered decoding, or `streaming.att_context_size` for cache-aware + models). The generated config is saved alongside the metrics log. + + Args: + cfg_path: Path to the NeMo config file. + metrics_log: Path to the output metrics log file; the generated config is saved in the same + directory. + src_lang: Source language code. + tgt_lang: Target language code. + overrides: List of "key=value" strings to override config fields. + reference_manifest: Optional path to a manifest with reference text, for WER calculation. + output_manifest: Optional path to write a NeMo-style prediction manifest. + + Returns: + Path to the generated config file with the required fields added. + """ + cfg = OmegaConf.load(cfg_path) + + if overrides: + logging.info("Applying command-line overrides:") + try: + override_conf = OmegaConf.from_dotlist(overrides) + cfg = OmegaConf.merge(cfg, override_conf) + for ov in overrides: + logging.info(f" {ov}") + except Exception as e: + logging.error(f" Error applying overrides {overrides}: {e}") + + if src_lang is not None: + cfg.nmt.source_language = get_language_name(src_lang) + if tgt_lang is not None: + cfg.nmt.target_language = get_language_name(tgt_lang) + + if 'type' in cfg: + logging.info(f"Config already has 'type' field, using as-is: {cfg_path}") + return cfg_path + + logging.info(f"Adding simulstream fields to config: {cfg_path}") + + if 'streaming' in cfg and 'chunk_size' in cfg.streaming: + speech_chunk_size = cfg.streaming.chunk_size + logging.info(f" Using chunk size from config: {speech_chunk_size}s for buffered decoding") + elif 'streaming' in cfg and 'att_context_size' in cfg.streaming: + speech_chunk_size = (cfg.streaming.att_context_size[1] + 1) * 0.08 + logging.info(f" Using chunk size calculated from att_context_size: {speech_chunk_size}s") + else: + raise ValueError(f"No chunk_size or att_context_size found in config: {cfg_path}") + + simulstream_fields = OmegaConf.create( + { + 'type': 'nemo.collections.asr.inference.utils.simulstream_pipeline_adapter.NeMoStreamingPipelineAdapter', + 'speech_chunk_size': speech_chunk_size, + 'detokenizer_type': 'simuleval', + 'latency_unit': get_latency_unit(tgt_lang), + } + ) + if output_manifest: + simulstream_fields.output_manifest_file = str(Path(output_manifest).resolve()) + if reference_manifest: + simulstream_fields.reference_manifest = str(Path(reference_manifest).resolve()) + + # Merge with simulstream fields taking precedence over anything already in the NeMo config. + cfg = OmegaConf.merge(simulstream_fields, cfg) + + out_dir = Path(metrics_log).resolve().parent + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / (Path(cfg_path).stem + '_simulstream.yaml') + with open(out_path, 'w') as f: + OmegaConf.save(cfg, f) + + logging.info(f" Saved config: {out_path}") + return str(out_path) + + +def main(): + parser = argparse.ArgumentParser( + description='Run simulstream inference with a NeMo streaming ASR(+NMT) pipeline', + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument('--config', required=True, help='Path to NeMo config file (YAML)') + + audio_group = parser.add_mutually_exclusive_group(required=True) + audio_group.add_argument('--wav-list', help='Path to text file containing audio file paths (one per line)') + audio_group.add_argument('--manifest', help='Path to NeMo manifest file (JSONL, with audio_filepath field)') + + parser.add_argument('--src-lang', required=True, help='Source language code (e.g., "ru", "en")') + parser.add_argument('--tgt-lang', required=True, help='Target language code (e.g., "en", "es")') + parser.add_argument( + '--metrics-log', default='metrics.jsonl', help='Path to output metrics log file (default: metrics.jsonl)' + ) + parser.add_argument( + '--output-manifest', + default=None, + help='Path to output prediction manifest JSONL (contains pred_text/pred_translation)', + ) + args, unknown_args = parser.parse_known_args() + + # Unknown args of the form key=value are passed through as NeMo config overrides. + overrides = [] + for arg in unknown_args: + if arg.startswith("--"): + logging.warning(f"Unknown argument: {arg}") + elif "=" in arg: + overrides.append(arg) + else: + logging.warning(f"Ignoring unknown argument (expected key=value): {arg}") + + wav_list_path = args.wav_list + temp_wav_list = None + try: + if args.manifest: + logging.info(f"Loading audio paths from manifest: {args.manifest}") + audio_paths = load_manifest_audio_paths(args.manifest) + if not audio_paths: + raise RuntimeError(f"No audio files found in manifest: {args.manifest}") + _, temp_wav_list = tempfile.mkstemp(suffix='.txt', prefix='wav_list_') + with open(temp_wav_list, 'w') as f: + for path in audio_paths: + f.write(f"{path}\n") + wav_list_path = temp_wav_list + logging.info(f"Created temporary wav list: {temp_wav_list}") + + config_path = add_simulstream_fields( + args.config, + args.metrics_log, + args.src_lang, + args.tgt_lang, + overrides, + reference_manifest=args.manifest, + output_manifest=args.output_manifest, + ) + + simulstream_cmd = shutil.which('simulstream_inference') + if not simulstream_cmd: + raise RuntimeError( + "simulstream_inference not found in PATH. " + "Make sure simulstream is installed (`pip install simulstream`) and in your PATH." + ) + + cmd = [ + simulstream_cmd, + '--speech-processor-config', + config_path, + '--wav-list-file', + wav_list_path, + '--src-lang', + args.src_lang, + '--tgt-lang', + args.tgt_lang, + '--metrics-log-file', + args.metrics_log, + ] + + logging.info( + f"Running simulstream inference: config={args.config}, " + f"audio={args.manifest or args.wav_list}, " + f"src_lang={args.src_lang}, tgt_lang={args.tgt_lang}, metrics_log={args.metrics_log}" + ) + + subprocess.run(cmd, check=True) + finally: + if temp_wav_list: + try: + Path(temp_wav_list).unlink() + logging.info(f"Cleaned up temporary wav list: {temp_wav_list}") + except OSError as e: + logging.warning(f"Failed to clean up temporary wav list {temp_wav_list}: {e}") + + +if __name__ == '__main__': + main() diff --git a/examples/asr/conf/asr_streaming_inference/buffered_ctc.yaml b/examples/asr/conf/asr_streaming_inference/buffered_ctc.yaml index 00500051125a..2e28d9637443 100644 --- a/examples/asr/conf/asr_streaming_inference/buffered_ctc.yaml +++ b/examples/asr/conf/asr_streaming_inference/buffered_ctc.yaml @@ -36,6 +36,8 @@ nmt: llm_params: # See https://docs.vllm.ai/en/v0.8.1/api/offline_inference/llm.html for more details dtype: "auto" # Compute precision seed: 42 # The seed to initialize the random number generator for sampling + gpu_memory_utilization: 0.85 # Sets what fraction of each GPU’s memory vLLM is allowed to reserve and use. Use smaller values for reducing GPU memory usage. + max_model_len: 4096 # Maximum length of the model tokens. sampling_params: # See https://docs.vllm.ai/en/v0.6.4/dev/sampling_params.html for more details max_tokens: 100 # Maximum number of tokens to generate with LLM temperature: 0.0 # LLM sampling temperature, default for translation is 0 (greedy) diff --git a/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml b/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml index ef3f5f776d28..75a446c71015 100644 --- a/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml +++ b/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml @@ -2,7 +2,7 @@ # ASR Configuration # ================================ asr: - model_name: nvidia/parakeet-rnnt-1.1b # Pre-trained RNNT/hybrid model from NGC/HuggingFace or local .nemo file path + model_name: nvidia/parakeet-tdt-0.6b-v2 # Pre-trained RNNT/hybrid model from NGC/HuggingFace or local .nemo file path device: cuda # Device for inference: 'cuda' or 'cpu' device_id: 0 # GPU device ID compute_dtype: bfloat16 # Compute precision: 'bfloat16' for Ampere+, 'float16' for older GPUs, or 'float32' @@ -59,6 +59,8 @@ nmt: llm_params: # See https://docs.vllm.ai/en/v0.8.1/api/offline_inference/llm.html for more details dtype: "auto" # Compute precision seed: 42 # The seed to initialize the random number generator for sampling + gpu_memory_utilization: 0.85 # Sets what fraction of each GPU’s memory vLLM is allowed to reserve and use. Use smaller values for reducing GPU memory usage. + max_model_len: 4096 # Maximum length of the model tokens. sampling_params: # See https://docs.vllm.ai/en/v0.6.4/dev/sampling_params.html for more details max_tokens: 100 # Maximum number of tokens to generate with LLM temperature: 0.0 # LLM sampling temperature, default for translation is 0 (greedy) diff --git a/examples/asr/conf/asr_streaming_inference/buffered_salm.yaml b/examples/asr/conf/asr_streaming_inference/buffered_salm.yaml index d2c46701dc1f..492a9c4c6a5d 100644 --- a/examples/asr/conf/asr_streaming_inference/buffered_salm.yaml +++ b/examples/asr/conf/asr_streaming_inference/buffered_salm.yaml @@ -36,13 +36,14 @@ nmt: llm_params: # See https://docs.vllm.ai/en/v0.8.1/api/offline_inference/llm.html for more details dtype: "auto" # Compute precision seed: 42 # The seed to initialize the random number generator for sampling + gpu_memory_utilization: 0.85 # Sets what fraction of each GPU’s memory vLLM is allowed to reserve and use. Use smaller values for reducing GPU memory usage. + max_model_len: 4096 # Maximum length of the model tokens. sampling_params: # See https://docs.vllm.ai/en/v0.6.4/dev/sampling_params.html for more details max_tokens: 100 # Maximum number of tokens to generate with LLM temperature: 0.0 # LLM sampling temperature, default for translation is 0 (greedy) top_p: 0.9 # The cumulative probability threshold for nucleus sampling seed: 42 # The seed to initialize the random number generator for sampling - # ======================== # Endpointing settings # ======================== diff --git a/examples/asr/conf/asr_streaming_inference/cache_aware_ctc.yaml b/examples/asr/conf/asr_streaming_inference/cache_aware_ctc.yaml index 90b64e149683..3b1dbfe0b052 100644 --- a/examples/asr/conf/asr_streaming_inference/cache_aware_ctc.yaml +++ b/examples/asr/conf/asr_streaming_inference/cache_aware_ctc.yaml @@ -36,6 +36,8 @@ nmt: llm_params: # See https://docs.vllm.ai/en/v0.8.1/api/offline_inference/llm.html for more details dtype: "auto" # Compute precision seed: 42 # The seed to initialize the random number generator for sampling + gpu_memory_utilization: 0.85 # Sets what fraction of each GPU’s memory vLLM is allowed to reserve and use. Use smaller values for reducing GPU memory usage. + max_model_len: 4096 # Maximum length of the model tokens. sampling_params: # See https://docs.vllm.ai/en/v0.6.4/dev/sampling_params.html for more details max_tokens: 100 # Maximum number of tokens to generate with LLM temperature: 0.0 # LLM sampling temperature, default for translation is 0 (greedy) diff --git a/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml b/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml index cfa62b225506..1b9373e89a40 100644 --- a/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml +++ b/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml @@ -74,6 +74,8 @@ nmt: llm_params: # See https://docs.vllm.ai/en/v0.8.1/api/offline_inference/llm.html for more details dtype: "auto" # Compute precision seed: 42 # The seed to initialize the random number generator for sampling + gpu_memory_utilization: 0.85 # Sets what fraction of each GPU’s memory vLLM is allowed to reserve and use. Use smaller values for reducing GPU memory usage. + max_model_len: 4096 # Maximum length of the model input tokens. sampling_params: # See https://docs.vllm.ai/en/v0.6.4/dev/sampling_params.html for more details max_tokens: 100 # Maximum number of tokens to generate with LLM temperature: 0.0 # LLM sampling temperature, default for translation is 0 (greedy) diff --git a/nemo/collections/asr/inference/__init__.py b/nemo/collections/asr/inference/__init__.py index 341a77c5bc66..af21e09aebe1 100644 --- a/nemo/collections/asr/inference/__init__.py +++ b/nemo/collections/asr/inference/__init__.py @@ -11,3 +11,17 @@ # 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. + +# SimulStream integration is an optional dependency (`pip install simulstream`); the adapter module +# guards its own `import simulstream` internally, so importing it here always succeeds. +from nemo.collections.asr.inference.utils.simulstream_manifest_utils import ( + load_manifest_audio_paths, + manifest_to_audio_definition, +) +from nemo.collections.asr.inference.utils.simulstream_pipeline_adapter import NeMoStreamingPipelineAdapter + +__all__ = [ + 'NeMoStreamingPipelineAdapter', + 'load_manifest_audio_paths', + 'manifest_to_audio_definition', +] diff --git a/nemo/collections/asr/inference/nmt/llm_translator.py b/nemo/collections/asr/inference/nmt/llm_translator.py index 34c50b17eef9..8bf13cd0b86a 100644 --- a/nemo/collections/asr/inference/nmt/llm_translator.py +++ b/nemo/collections/asr/inference/nmt/llm_translator.py @@ -15,11 +15,18 @@ import os import string +from typing import Final import torch +from huggingface_hub import snapshot_download +from huggingface_hub.utils import LocalEntryNotFoundError from omegaconf import DictConfig, OmegaConf -from nemo.collections.asr.inference.nmt.prompts import EuroLLMTranslatorPromptTemplate, PromptTemplate +from nemo.collections.asr.inference.nmt.prompts import ( + EuroLLMTranslatorPromptTemplate, + PromptTemplate, + QwenReasoningTranslatorPromptTemplate, +) try: from vllm import LLM, SamplingParams @@ -28,9 +35,25 @@ from nemo.utils import logging -EURO_LLM_INSTRUCT_SMALL = "utter-project/EuroLLM-1.7B-Instruct" -EURO_LLM_INSTRUCT_LARGE = "utter-project/EuroLLM-9B-Instruct" -SUPPORTED_TRANSLATION_MODELS = [EURO_LLM_INSTRUCT_SMALL, EURO_LLM_INSTRUCT_LARGE] +TRANSLATION_MODELS_BY_SERIES: Final[dict[str, tuple[str, ...]]] = { + "eurollm": ( + "utter-project/EuroLLM-1.7B-Instruct", + "utter-project/EuroLLM-9B-Instruct", + ), + "qwen3": ( + "Qwen/Qwen3-4B-Instruct-2507", + "Qwen/Qwen3-8B", + ), + "qwen3.5": ( + "Qwen/Qwen3.5-4B", + "Qwen/Qwen3.5-9B", + "Qwen/Qwen3.5-27B", + "Qwen/Qwen3.5-35B-A3B", + ), +} +SUPPORTED_TRANSLATION_MODELS: Final[tuple[str, ...]] = tuple( + model for series_models in TRANSLATION_MODELS_BY_SERIES.values() for model in series_models +) class LLMTranslator: @@ -145,16 +168,28 @@ def get_prompt_template(model_name: str) -> PromptTemplate: Raises: ValueError: if model is not supported for translation """ - if model_name in [EURO_LLM_INSTRUCT_SMALL, EURO_LLM_INSTRUCT_LARGE]: + if model_name not in SUPPORTED_TRANSLATION_MODELS: + raise ValueError( + f"Model {model_name} is not supported for translation. Supported models are: {SUPPORTED_TRANSLATION_MODELS}" + ) + + if model_name in TRANSLATION_MODELS_BY_SERIES["eurollm"]: return EuroLLMTranslatorPromptTemplate - raise ValueError( - f"Model {model_name} is not supported for translation. Supported models are: {SUPPORTED_TRANSLATION_MODELS}" - ) + # Instruct qwen model template is similar to EuroLLM, so we use the same prompt template + if model_name == "Qwen/Qwen3-4B-Instruct-2507": + return EuroLLMTranslatorPromptTemplate + + if ( + model_name in TRANSLATION_MODELS_BY_SERIES["qwen3.5"] + or model_name in TRANSLATION_MODELS_BY_SERIES["qwen3"] + ): + return QwenReasoningTranslatorPromptTemplate def load_model(self, llm_params: dict) -> LLM: """ Load NMT model in vLLM format. + If the model is not found in the local cache, it will be downloaded from the HuggingFace model hub. Args: llm_params: (dict) parameters for the LLM model Returns: @@ -164,10 +199,38 @@ def load_model(self, llm_params: dict) -> LLM: """ try: os.environ["CUDA_VISIBLE_DEVICES"] = str(self.device_id) - model = LLM(model=self.model_name, **llm_params) + local_path = self._get_local_model_path(self.model_name) + if local_path is not None and os.path.exists(local_path): + logging.info(f"Loading LLM from local cache path: {local_path}") + model_name = local_path + else: + logging.info(f"Loading LLM from model name: {self.model_name}") + model_name = self.model_name + model = LLM(model=model_name, **llm_params) return model except Exception as e: - raise RuntimeError(f"Model loading failed: {str(e)}") from e + raise RuntimeError(f"Model loading failed: {str(e)}") + + def _get_local_model_path(self, repo_id): + """ + Get local model path from HuggingFace model hub. + Args: + repo_id: (str) repository ID of the model + Returns: + local_path: (str) local path of the model + Raises: + LocalEntryNotFoundError: If model is not found in the local cache + """ + try: + return snapshot_download( + repo_id=repo_id, + local_files_only=True, + ) + except LocalEntryNotFoundError: + logging.warning( + f"Model {repo_id} is not found in the local cache. Downloading from HuggingFace model hub." + ) + return None def translate_batch( self, @@ -202,7 +265,7 @@ def translate_batch( for tgt_prefix, output in zip(prefixes, outputs): output_text = output.outputs[0].text output_text = self.prompt_template.extract(output_text) - translations.append(f"{tgt_prefix}{output_text}") + translations.append(f"{tgt_prefix} {output_text.strip()}") return translations def translate( diff --git a/nemo/collections/asr/inference/nmt/prompts.py b/nemo/collections/asr/inference/nmt/prompts.py index c2f62a5ebae8..b4f185f7e25e 100644 --- a/nemo/collections/asr/inference/nmt/prompts.py +++ b/nemo/collections/asr/inference/nmt/prompts.py @@ -94,3 +94,116 @@ def extract(cls, response: str) -> str: str: The text before the first newline. """ return response.split('\n')[0] + + +class QwenReasoningTranslatorPromptTemplate(PromptTemplate): + """ + Chat-style prompt template for Qwen Reasoning model to perform translation. + Thinking is disabled by appending the empty think block (\\n\\n\\n\\n) + after <|im_start|>assistant\\n, matching the tokenizer.apply_chat_template(..., enable_thinking=False) behavior. + + The system and user prompts reproduce those used in the NVIDIA NeMo team's IWSLT 2026 submission: + Grigoryan, Bataev, Andrusenko, et al. (2026), "NeMo@IWSLT 2026: Cascaded System for Simultaneous + Speech Translation" (https://aclanthology.org/2026.iwslt-1.23.pdf). + """ + + SYSTEM_MESSAGE = ( + "You are a professional machine translation assistant.\n" + "Translate the input text into the target language.\n" + "- Output only the translation.\n" + "- Do not complete or extend the text.\n" + "- The input may be incomplete; preserve incompleteness.\n" + "- Do not infer missing content.\n" + "- Stop immediately after translating.\n" + "- Preserve named entities, numbers, punctuation, and formatting." + ) + + # Empty think block appended so model goes straight to answer (enable_thinking=False behavior) + _THINK_DISABLED_SUFFIX = "\n\n\n\n" + + USER_CONTENT_TEMPLATE = ( + "Translate the following {src_lang} source text to {tgt_lang}:\n" "{src_lang}: {src_text}\n" "{tgt_lang}: " + ) + + @classmethod + def format( + cls, + src_lang: str, + tgt_lang: str, + src_prefix: str, + tgt_prefix: str, + src_context: str = "", + tgt_context: str = "", + ) -> str: + """ + Generate a translation prompt in Qwen3 chat format (thinking disabled). + Args: + src_lang, tgt_lang, src_prefix, tgt_prefix, src_context, tgt_context: same as other templates. + Returns: + str: Formatted prompt string. + """ + src_text = f"{src_context} {src_prefix}" + tgt_text = f"{tgt_context} {tgt_prefix}" + src_text = re.sub(r"\s+", " ", src_text).strip() + tgt_text = re.sub(r"\s+", " ", tgt_text).strip() + user_content = cls.USER_CONTENT_TEMPLATE.format( + src_lang=src_lang, tgt_lang=tgt_lang, src_text=src_text, tgt_text=tgt_text + ) + assistant_text = f"{cls._THINK_DISABLED_SUFFIX}{tgt_text}" + + system_block = f"<|im_start|>system\n{cls.SYSTEM_MESSAGE}<|im_end|>\n" + return ( + system_block + + f"<|im_start|>user\n{user_content}<|im_end|>\n" + + f"<|im_start|>assistant\n{assistant_text}" + ) + + @classmethod + def messages( + cls, + src_lang: str, + tgt_lang: str, + src_prefix: str, + tgt_prefix: str, + src_context: str = "", + tgt_context: str = "", + ): + """ + Return chat messages for tokenizer.apply_chat_template(). + System message instructs the model not to use (thinking disabled). + """ + src_text = re.sub(r"\s+", " ", f"{src_context} {src_prefix}".strip()).strip() + tgt_text = re.sub(r"\s+", " ", f"{tgt_context} {tgt_prefix}".strip()).strip() + user_content = cls.USER_CONTENT_TEMPLATE.format( + src_lang=src_lang, tgt_lang=tgt_lang, src_text=src_text, tgt_text=tgt_text + ) + return [ + {"role": "system", "content": cls.SYSTEM_MESSAGE}, + {"role": "user", "content": user_content}, + {"role": "assistant", "content": tgt_text}, + ] + + @classmethod + def extract(cls, response: str) -> str: + """ + Extract the translation from the model response. Strips any think block + (...) so only the actual translation is returned (thinking disabled + at decode time). Falls back to first line if no think block is present. + """ + response = response.strip() + if "" in response: + response = response.split("")[-1].strip() + if "" in response: + response = response.split("")[-1].strip() + if not response: + return "" + + # Remove any trailing punctuation to reduce the risk of hallucination + response = response.removesuffix("...") + if not response: + return "" + + parts = response.rsplit(maxsplit=1) + parts[-1] = parts[-1].replace("...", "") + response = " ".join(parts) + return response diff --git a/nemo/collections/asr/inference/pipelines/base_pipeline.py b/nemo/collections/asr/inference/pipelines/base_pipeline.py index 178afaaf3dec..6488e2ca6df7 100644 --- a/nemo/collections/asr/inference/pipelines/base_pipeline.py +++ b/nemo/collections/asr/inference/pipelines/base_pipeline.py @@ -73,6 +73,10 @@ class TranscribeStepOutput: # Current step transcript/translation is the transcript/translation generated from the current frame current_step_transcript: str = "" current_step_translation: str = "" + # Partial translation as of the previous step, before this step's translate_step() ran. + # Consumers (e.g. incremental output adapters) use this to diff against partial_translation + # and compute what was added/removed since the last step. + previous_partial_translation: str = "" @classmethod def from_state(cls, state: StreamingState, request: Request, sep: str = ' ') -> 'TranscribeStepOutput': @@ -104,6 +108,7 @@ def from_state(cls, state: StreamingState, request: Request, sep: str = ' ') -> final_segments=final_segments, partial_transcript=state.partial_transcript, current_step_transcript=state.current_step_transcript, + previous_partial_translation=state.previous_translation_info[0], ) def __str__(self) -> str: @@ -211,6 +216,9 @@ def translate_step(self, states: list[StreamingState], step_outputs: list[Transc final = step_output.final_transcript partial = step_output.partial_transcript if not (final.strip() or partial.strip()): + # No new transcript to translate this step: keep the previous partial translation + step_output.previous_partial_translation = state.previous_translation_info[0] + step_output.partial_translation = state.previous_translation_info[0] continue transcript = final or partial @@ -240,6 +248,8 @@ def translate_step(self, states: list[StreamingState], step_outputs: list[Transc for (state, step_output), translation, new_prefix, prev_prefix, is_final in zip( states_to_translate, translations, new_prefixes, current_prefixes, final_transcript_mask ): + step_output.previous_partial_translation = state.previous_translation_info[0] + if is_final: step_output.final_translation = translation step_output.partial_translation = "" diff --git a/nemo/collections/asr/inference/utils/simulstream_manifest_utils.py b/nemo/collections/asr/inference/utils/simulstream_manifest_utils.py new file mode 100644 index 000000000000..cba184b0091d --- /dev/null +++ b/nemo/collections/asr/inference/utils/simulstream_manifest_utils.py @@ -0,0 +1,109 @@ +# 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. + +""" +Utilities for using NeMo manifest files with simulstream evaluation. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import yaml + +from nemo.utils import logging + + +def load_manifest_audio_paths(manifest_path: str) -> list[str]: + """ + Load audio file paths from a NeMo manifest file. + + Args: + manifest_path: Path to NeMo manifest JSONL file + + Returns: + List of audio file paths + """ + audio_paths = [] + manifest_dir = Path(manifest_path).parent + + with open(manifest_path, 'r', encoding='utf-8') as f: + for line_num, line in enumerate(f, 1): + line = line.strip() + if not line: + continue + try: + data = json.loads(line) + except json.JSONDecodeError as e: + logging.warning(f"Failed to parse line {line_num} in manifest: {e}") + continue + audio_path = data.get('audio_filepath', data.get('audio_file')) + if audio_path: + audio_path = Path(audio_path) + if not audio_path.is_absolute(): + audio_path = manifest_dir / audio_path + audio_paths.append(str(audio_path.resolve())) + + logging.info(f"Loaded {len(audio_paths)} audio files from manifest: {manifest_path}") + return audio_paths + + +def manifest_to_audio_definition(manifest_path: str, output_path: str) -> tuple[Path, Path, Path]: + """ + Create simulstream audio definition YAML from a NeMo manifest, along with plaintext + reference/transcript files. This is needed for simulstream's score/latency metrics evaluation. + + Args: + manifest_path: Path to NeMo manifest file. + output_path: Directory to write the generated files into. + + Returns: + Tuple of (audio_definitions.yaml, references.txt, transcripts.txt) paths. + """ + audio_defs = [] + references = [] + transcripts = [] + + with open(manifest_path, 'r', encoding='utf-8') as f: + for line in f: + if not line.strip(): + continue + data = json.loads(line.strip()) + + audio_path = data['audio_filepath'] + duration = data.get('duration', 0.0) + audio_defs.append({'wav': audio_path, 'offset': 0.0, 'duration': float(duration) if duration else 0.0}) + + transcripts.append(data.get('text', '')) + # Prefer 'target_text', falling back to 'answer' (common NeMo AST manifest field). + references.append(data.get('target_text', data.get('answer', ''))) + + output_dir = Path(output_path) + output_dir.mkdir(parents=True, exist_ok=True) + + audio_def_file = output_dir / 'audio_definitions.yaml' + with open(audio_def_file, 'w', encoding='utf-8') as f: + yaml.dump(audio_defs, f, default_flow_style=False, allow_unicode=True) + + refs_file = output_dir / 'references.txt' + with open(refs_file, 'w', encoding='utf-8') as f: + f.writelines(ref + '\n' for ref in references) + + trans_file = output_dir / 'transcripts.txt' + with open(trans_file, 'w', encoding='utf-8') as f: + f.writelines(trans + '\n' for trans in transcripts) + + logging.info(f"Created simulstream audio definition files in: {output_dir}") + return audio_def_file, refs_file, trans_file diff --git a/nemo/collections/asr/inference/utils/simulstream_pipeline_adapter.py b/nemo/collections/asr/inference/utils/simulstream_pipeline_adapter.py new file mode 100644 index 000000000000..e71392f1a540 --- /dev/null +++ b/nemo/collections/asr/inference/utils/simulstream_pipeline_adapter.py @@ -0,0 +1,584 @@ +#!/usr/bin/env python3 +# 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. + +""" +Adapter to use NeMo's native streaming pipelines with simulstream evaluation. + +This adapter properly interfaces with NeMo's internal streaming API (transcribe_step) +rather than duplicating chunking/buffering logic. NeMo handles all buffering internally. + +Key Insight: + NeMo's pipelines already have complete streaming infrastructure: + - Frame creation and buffering logic (BufferedRNNTPipeline / CacheAwareRNNTPipeline) + - State management (StreamingState) + - Translation integration (LLMTranslator) + + We just need to: + 1. Create Frame requests from audio chunks + 2. Call pipeline.transcribe_step() + 3. Convert TranscribeStepOutput -> IncrementalOutput +""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import List, Optional + +import numpy as np +import torch +from omegaconf import OmegaConf + +from nemo.collections.asr.parts.context_biasing.biasing_multi_model import BiasingRequestItemConfig +from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import BoostingTreeModelConfig +from nemo.collections.asr.parts.utils.eval_utils import cal_write_wer +from nemo.utils import logging + +try: + from simulstream.server.speech_processors import SpeechProcessor + from simulstream.server.speech_processors.incremental_output import IncrementalOutput + + SIMULSTREAM_AVAILABLE = True +except ImportError: + SIMULSTREAM_AVAILABLE = False + SpeechProcessor = object + + +class NeMoStreamingPipelineAdapter(SpeechProcessor): + """ + Adapter to use NeMo's streaming pipelines with simulstream evaluation. + + Architecture: + audio_chunk -> Frame -> pipeline.transcribe_step() -> TranscribeStepOutput -> IncrementalOutput + + The pipeline internally handles: + - Buffering (cache-aware or buffered mode) + - Feature extraction + - ASR decoding (CTC/RNN-T) + - Translation (optional, via LLMTranslator) + - State management per stream + """ + + pipeline = None # Class-level pipeline (shared across instances) + output_manifest_path: Optional[str] = None + wav_names: list[str] = [] + per_stream_boosting_requests: list[BiasingRequestItemConfig] | None = None + + def __init__(self, config: SimpleNamespace): + """ + Initialize adapter. + + Args: + config: Configuration from simulstream (SimpleNamespace). Will be converted to an + OmegaConf DictConfig for NeMo in `load_model`. + """ + if not SIMULSTREAM_AVAILABLE: + raise ImportError("simulstream is required. Install with: pip install simulstream") + + super().__init__(config) + + self.stream_id = 0 + self._reset_stream_state() + + self.latency_unit = getattr(config, 'latency_unit', 'word') + if isinstance(self.latency_unit, str): + self.latency_unit = self.latency_unit.lower() + if self.latency_unit not in ("word", "char"): + logging.warning(f"Unsupported latency_unit='{self.latency_unit}', defaulting to 'word'") + self.latency_unit = "word" + + # Language settings (set at runtime by simulstream via set_source_language/set_target_language) + self.src_lang = None + self.tgt_lang = None + + @classmethod + def load_model(cls, config: SimpleNamespace): + """ + Load the NeMo pipeline once (class-level, shared across all stream instances). + + Args: + config: Configuration from simulstream. + """ + if cls.pipeline is not None: + return # Already loaded + + import atexit + + from nemo.collections.asr.inference.factory.pipeline_builder import PipelineBuilder + + # SimulStream configs are SimpleNamespace objects; NeMo expects an OmegaConf DictConfig. + cfg = OmegaConf.create(cls._namespace_to_dict(config)) + + # This adapter always sends raw-audio Frame requests (see process_chunk), so the pipeline + # must be built to match: `streaming.request_type` controls whether BufferedRNNTPipeline + # builds a Frame- or feature_buffer-based bufferer, and mismatching it would silently + # produce unpadded/wrong features. Force it here instead of relying on every config. + if cfg.get('streaming', {}).get('request_type', 'frame') != 'frame': + logging.warning( + f"Overriding streaming.request_type='{cfg.streaming.request_type}' to 'frame' " + f"({type(cls).__name__} only supports frame requests)." + ) + cfg.streaming.request_type = 'frame' + cls.cfg = cfg + + cls.pipeline = PipelineBuilder.build_pipeline(cfg) + cls.pipeline.open_session() + + cls.detailed_log_path = getattr(config, "detailed_log_path", None) + + # Output manifest path (optional, but derived from metrics_log_file when not set explicitly). + cls.output_manifest_path = getattr(config, 'output_manifest_file', None) or getattr( + config, 'output_manifest', None + ) + if cls.output_manifest_path is None: + metrics_log_file = getattr(config, 'metrics_log_file', None) + if metrics_log_file: + metrics_path = Path(metrics_log_file) + cls.output_manifest_path = str(metrics_path.parent / f"{metrics_path.stem}_pred_manifest.jsonl") + + if cls.output_manifest_path: + Path(cls.output_manifest_path).write_text("", encoding="utf-8") # truncate at start of run + logging.info(f"Prediction manifest output: {cls.output_manifest_path}") + cls._wer_calculated = False + + cls.wav_names = [] + wav_list_file = getattr(config, 'wav_list_file', None) + if wav_list_file and Path(wav_list_file).exists(): + with open(wav_list_file, 'r', encoding='utf-8') as f: + cls.wav_names = [line.strip() for line in f if line.strip()] + + cls._load_reference_manifest(config) + + # vLLM (used by LLMTranslator) needs to be shut down explicitly, otherwise it can hang or + # print noisy errors on process exit. + atexit.register(cls.cleanup_model) + + logging.info(f"Loaded NeMo pipeline: {type(cls.pipeline).__name__}") + logging.info(f" ASR model: {cfg.asr.model_name}") + if cfg.get('enable_nmt', False): + logging.info(f" NMT model: {cfg.nmt.model_name}") + logging.info(f" Translation: {cfg.nmt.source_language} -> {cfg.nmt.target_language}") + + if cfg.get("per_stream_boosting") and cfg.per_stream_boosting.get("phrases_file"): + boosting_model_alpha = cfg.per_stream_boosting.get("alpha", 1.0) + with open(cfg.per_stream_boosting.phrases_file, "r", encoding="utf-8") as f: + boosting_requests_raw = json.load(f) + cls.per_stream_boosting_requests = [ + BiasingRequestItemConfig( + BoostingTreeModelConfig(key_phrases_list=item["key_phrases_list"]), + boosting_model_alpha=boosting_model_alpha, + ) + for item in boosting_requests_raw + ] + logging.info( + f"Per-stream boosting enabled with weight {boosting_model_alpha:.2g}, " + f"expected {len(cls.per_stream_boosting_requests)} ordered streams" + ) + else: + logging.info( + "Per-stream boosting disabled; to enable, " + "specify `per_stream_boosting.phrases_file` and `per_stream_boosting.alpha`" + ) + + @classmethod + def _load_reference_manifest(cls, config: SimpleNamespace) -> None: + """Load optional input manifest to copy reference text fields and enable WER calculation.""" + cls.reference_manifest_by_audio = {} + cls.reference_manifest_by_basename = {} + cls.reference_manifest_items_ordered = [] + + manifest_path = None + for key in ( + "manifest", + "manifest_file", + "input_manifest", + "input_manifest_file", + "reference_manifest", + ): + value = getattr(config, key, None) + if value: + manifest_path = value + break + + if not manifest_path: + return + + manifest_path = str(manifest_path) + if not Path(manifest_path).exists(): + logging.warning(f"Reference manifest path not found: {manifest_path}") + return + + manifest_dir = Path(manifest_path).parent + loaded = 0 + with open(manifest_path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + item = json.loads(line) + except json.JSONDecodeError: + continue + audio = item.get("audio_filepath", "") + if not audio: + continue + audio_path = Path(audio) + if not audio_path.is_absolute(): + audio_path = manifest_dir / audio_path + audio_abs = str(audio_path.resolve()) + cls.reference_manifest_by_audio[audio_abs] = item + cls.reference_manifest_by_basename[audio_path.name] = item + cls.reference_manifest_items_ordered.append(item) + loaded += 1 + + logging.info(f"Loaded reference manifest entries: {loaded}") + + @staticmethod + def _namespace_to_dict(obj): + """Recursively convert SimpleNamespace to dict.""" + if isinstance(obj, SimpleNamespace): + return {k: NeMoStreamingPipelineAdapter._namespace_to_dict(v) for k, v in vars(obj).items()} + elif isinstance(obj, dict): + return {k: NeMoStreamingPipelineAdapter._namespace_to_dict(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [NeMoStreamingPipelineAdapter._namespace_to_dict(item) for item in obj] + return obj + + def set_source_language(self, language: str) -> None: + """Set source language (simulstream interface).""" + self.src_lang = language + + def set_target_language(self, language: str) -> None: + """Set target language (simulstream interface).""" + self.tgt_lang = language + + def process_chunk(self, audio: np.ndarray) -> IncrementalOutput: + """ + Process an audio chunk using NeMo's native streaming API. + + Creates a Frame request and calls pipeline.transcribe_step(), which internally + handles all buffering, feature extraction, and decoding. + + NOTE: works only with batch size 1 (so does SimulStream). + + Args: + audio: Audio chunk (numpy array, float32, mono, 16kHz) + + Returns: + IncrementalOutput: Streaming results (partial/final ASR + translation) + """ + from nemo.collections.asr.inference.streaming.framing.request import Frame + from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions + + if audio.ndim > 1: + raise ValueError("Simulstream processes only one audio at a time (batch size 1).") + + expected_chunk_size = int(16000 * self.speech_chunk_size) + audio_length = len(audio) + if audio_length < expected_chunk_size: + audio = np.concatenate([audio, np.zeros(expected_chunk_size - audio_length)]) + audio_tensor = torch.from_numpy(audio).float().to(self.pipeline.device) + + if self.is_first_chunk and self.per_stream_boosting_requests is not None: + biasing_cfg = self.per_stream_boosting_requests[self.stream_id] + else: + biasing_cfg = None + + # simulstream doesn't tell us whether a chunk is the last one; is_last is always False here, + # and any leftover right-context is handled by end_of_stream()/return_tail_result. The + # pipeline's own internal bufferer accumulates the left/right padding sliding window + # per stream from these raw-audio Frames (see __init__ for why "frame" is the only + # supported request type). + request = Frame( + stream_id=self.stream_id, + samples=audio_tensor, + is_first=self.is_first_chunk, + is_last=False, + length=audio_length, + options=ASRRequestOptions(biasing_cfg=biasing_cfg) if self.is_first_chunk else None, + ) + + # This internally handles: buffering -> encoding -> decoding -> translation + step_outputs = self.pipeline.transcribe_step([request]) + step_output = step_outputs[0] + + # Snapshot the previous accumulated transcript before updating it below, so + # _convert_to_incremental_output can diff against it when NMT is disabled. + previous_transcript = self._final_transcript_acc + self._last_partial_transcript + + # Track final/latest-partial outputs to write a NeMo-style prediction manifest line. + self._final_transcript_acc += step_output.final_transcript or "" + self._final_translation_acc += step_output.final_translation or "" + self._last_partial_transcript = step_output.partial_transcript or "" + if step_output.final_translation: + self._last_partial_translation = step_output.final_translation + elif step_output.partial_translation: + self._last_partial_translation = step_output.partial_translation + + result = self._convert_to_incremental_output(step_output, previous_transcript) + + self.is_first_chunk = False + + if self.detailed_log_path is not None: + with open(self.detailed_log_path, "a", encoding="utf-8") as f: + print( + json.dumps( + { + "final_transcript": step_output.final_transcript, + "partial_transcript": step_output.partial_transcript, + "final_translation": step_output.final_translation, + "partial_translation": step_output.partial_translation, + "new_tokens": result.new_tokens, + "new_string": result.new_string, + "deleted_tokens": result.deleted_tokens, + "deleted_string": result.deleted_string, + } + ), + file=f, + ) + + return result + + def _convert_to_incremental_output(self, step_output, previous_transcript: str = "") -> IncrementalOutput: + """ + Convert NeMo's TranscribeStepOutput to simulstream's IncrementalOutput. + + Computes generated/deleted tokens by diffing the previous and current partial (or final) + output, tokenized according to `latency_unit` (word-split, or per-character for languages + like Chinese). When NMT is enabled, the translation is diffed (since re-translation from + the current transcript prefix can revise earlier text); otherwise the ASR transcript is + diffed directly. + + Args: + step_output: NeMo's TranscribeStepOutput for the current chunk. + previous_transcript: Full accumulated transcript (final + partial) before this step, + used as the diff baseline when NMT is disabled. + + Returns: + IncrementalOutput: Simulstream format with generated/deleted token lists. + """ + if self.pipeline.nmt_enabled: + prev_partial = step_output.previous_partial_translation + if step_output.final_translation: + current_partial = step_output.final_translation + elif step_output.partial_translation: + current_partial = step_output.partial_translation + else: + current_partial = "" + else: + prev_partial = previous_transcript + current_partial = self._final_transcript_acc + self._last_partial_transcript + + prev_tokens = self._tokenize_text(prev_partial) + curr_tokens = self._tokenize_text(current_partial) + + common_prefix_len = 0 + for i in range(min(len(prev_tokens), len(curr_tokens))): + if prev_tokens[i] == curr_tokens[i]: + common_prefix_len += 1 + else: + break + + deleted_tokens = prev_tokens[common_prefix_len:] + generated_tokens = curr_tokens[common_prefix_len:] + + return IncrementalOutput( + new_tokens=generated_tokens, + new_string=self._join_tokens(generated_tokens), + deleted_tokens=deleted_tokens, + deleted_string=self._join_tokens(deleted_tokens), + ) + + def end_of_stream(self) -> IncrementalOutput: + """ + Called at the end of the audio stream to finalize output. + + The last chunk was already processed with is_last=False in process_chunk() (simulstream + doesn't signal which chunk is last), so this only finalizes stream state / writes the + prediction manifest line and emits an empty incremental output. Required by the + SpeechProcessor interface. + """ + pred_text = (self._final_transcript_acc + self._last_partial_transcript).strip() + pred_translation = (self._final_translation_acc + self._last_partial_translation).strip() + self._write_prediction_manifest_line(pred_text, pred_translation) + + self.pipeline.delete_state(self.stream_id) + return IncrementalOutput(new_tokens=[], new_string="", deleted_tokens=[], deleted_string="") + + def clear(self) -> None: + """ + Clear stream state and prepare for the next audio stream (simulstream interface). + """ + if not self.is_first_chunk: + self.end_of_stream() + + self.stream_id += 1 + self._reset_stream_state() + + def _reset_stream_state(self) -> None: + """Reset per-stream accumulators (used on init and between streams).""" + self.is_first_chunk = True + self._final_transcript_acc = "" + self._final_translation_acc = "" + self._last_partial_transcript = "" + self._last_partial_translation = "" + + def _write_prediction_manifest_line(self, pred_text: str, pred_translation: str) -> None: + """Write one NeMo-style manifest line with model predictions.""" + if not self.output_manifest_path: + return + + audio_filepath = "" + if self.stream_id < len(self.wav_names): + audio_filepath = self.wav_names[self.stream_id] + + reference_item = self._get_reference_item(audio_filepath) + if not audio_filepath and reference_item is not None: + audio_filepath = str(reference_item.get("audio_filepath", "") or "") + reference_text = "" + reference_translation = "" + if reference_item is not None: + reference_text = reference_item.get("text", "") + reference_translation = reference_item.get("answer", "") + + item = { + "audio_filepath": audio_filepath, + "text": reference_text, + "translation": reference_translation, + "pred_text": pred_text, + "pred_translation": pred_translation, + } + + with open(self.output_manifest_path, 'a', encoding='utf-8') as f: + f.write(json.dumps(item, ensure_ascii=False) + "\n") + + # Compute WER once, when the last stream's line is flushed. + if self.wav_names and self.stream_id == len(self.wav_names) - 1: + self._calculate_and_write_wer() + + def _get_reference_item(self, audio_filepath: str) -> Optional[dict]: + """Get reference manifest item by absolute path, basename, or stream order.""" + if not audio_filepath: + if self.stream_id < len(self.reference_manifest_items_ordered): + return self.reference_manifest_items_ordered[self.stream_id] + return None + try: + audio_abs = str(Path(audio_filepath).resolve()) + except Exception: + audio_abs = audio_filepath + item = self.reference_manifest_by_audio.get(audio_abs) + if item is not None: + return item + item = self.reference_manifest_by_basename.get(Path(audio_filepath).name) + if item is not None: + return item + if self.stream_id < len(self.reference_manifest_items_ordered): + return self.reference_manifest_items_ordered[self.stream_id] + return None + + @classmethod + def _calculate_and_write_wer(cls) -> None: + """Calculate WER from the output manifest and write summary artifacts.""" + if cls._wer_calculated or not cls.output_manifest_path: + return + + gt_text_attr_name = "text" + clean_groundtruth_text = False + langid = "en" + use_cer = False + ignore_capitalization = False + ignore_punctuation = False + + try: + if cls.cfg is not None and cls.cfg.get("metrics") and cls.cfg.metrics.get("asr"): + asr_cfg = cls.cfg.metrics.asr + gt_text_attr_name = asr_cfg.get("gt_text_attr_name", gt_text_attr_name) + clean_groundtruth_text = asr_cfg.get("clean_groundtruth_text", clean_groundtruth_text) + langid = asr_cfg.get("langid", langid) + use_cer = asr_cfg.get("use_cer", use_cer) + ignore_capitalization = asr_cfg.get("ignore_capitalization", ignore_capitalization) + ignore_punctuation = asr_cfg.get("ignore_punctuation", ignore_punctuation) + except Exception as e: + logging.warning(f"Failed to read ASR metric config, using defaults: {e}") + + try: + output_manifest_w_wer, total_res, _ = cal_write_wer( + pred_manifest=cls.output_manifest_path, + gt_text_attr_name=gt_text_attr_name, + pred_text_attr_name="pred_text", + output_filename=None, + clean_groundtruth_text=clean_groundtruth_text, + langid=langid, + use_cer=use_cer, + ignore_capitalization=ignore_capitalization, + ignore_punctuation=ignore_punctuation, + ) + + if output_manifest_w_wer: + metrics_summary_path = str(Path(cls.output_manifest_path).with_suffix(".wer.txt")) + with open(metrics_summary_path, "w", encoding="utf-8") as f: + f.write(str(total_res) + "\n") + logging.info(f"WER manifest: {output_manifest_w_wer}") + logging.info(f"WER summary: {metrics_summary_path}") + else: + logging.warning("WER calculation skipped because ground-truth text is unavailable in output manifest.") + except Exception as e: + logging.warning(f"Failed to calculate WER: {e}") + finally: + cls._wer_calculated = True + + def tokens_to_string(self, tokens: List[str]) -> str: + """Convert a token sequence into a human-readable string (SpeechProcessor interface).""" + return self._join_tokens(tokens) + + def _tokenize_text(self, text: Optional[str]) -> List[str]: + """Tokenize text according to the configured latency unit (word or char).""" + if not text: + return [] + text = text.replace("…", "") # keep token counts consistent with simulstream's own eval path + if self.latency_unit == "char": + return list(text.strip()) + return text.strip().split() + + def _join_tokens(self, tokens: List[str]) -> str: + """Join tokens according to the configured latency unit.""" + if not tokens: + return "" + if self.latency_unit == "char": + return "".join(tokens) + return " ".join(tokens) + + @classmethod + def cleanup_model(cls): + """ + Explicitly clean up vLLM (used by NMT) and release resources. Registered as an atexit + handler so the vLLM engine shuts down gracefully instead of erroring on process exit. + """ + if cls.pipeline is not None: + cls._calculate_and_write_wer() + if cls.pipeline is not None and cls.pipeline.nmt_model is not None: + try: + if hasattr(cls.pipeline.nmt_model, 'nmt_model'): + vllm_engine = cls.pipeline.nmt_model.nmt_model + if hasattr(vllm_engine, 'llm_engine'): + from vllm.distributed import destroy_model_parallel + + destroy_model_parallel() + del vllm_engine + cls.pipeline.nmt_model.nmt_model = None + logging.info("vLLM engine cleaned up") + except Exception as e: + logging.warning(f"Error during vLLM cleanup: {e}") diff --git a/nemo/collections/asr/parts/utils/eval_utils.py b/nemo/collections/asr/parts/utils/eval_utils.py index c7d23a214bbc..4607a237b3fb 100644 --- a/nemo/collections/asr/parts/utils/eval_utils.py +++ b/nemo/collections/asr/parts/utils/eval_utils.py @@ -233,7 +233,7 @@ def cal_write_wer( with open(output_manifest_w_wer, 'w') as fout: for sample in samples: - json.dump(sample, fout) + json.dump(sample, fout, ensure_ascii=False) fout.write('\n') fout.flush() @@ -319,7 +319,7 @@ def cal_write_text_metric( with open(output_manifest_w_wer, 'w') as fout: for sample in samples: - json.dump(sample, fout) + json.dump(sample, fout, ensure_ascii=False) fout.write('\n') fout.flush()