Add Riva Translate v2 support

@@ -28,6 +28,7 @@     PromptTemplate,     QwenReasoningTranslatorPromptTemplate,     RivaTranslatorPromptTemplate,+    RivaV2TranslatorPromptTemplate, )  try:@@ -37,6 +38,18 @@  from nemo.utils import logging +RIVA_V2_MODEL_NAME: Final[str] = "nvidia/Riva-Translate-4B-Instruct-v2"+CHINESE_LANGUAGE_NAMES: Final[frozenset[str]] = frozenset(+    {+        "chinese",+        "simplified chinese",+        "traditional chinese",+        "zh",+        "zh-cn",+        "zh-tw",+    }+)+ TRANSLATION_MODELS_BY_SERIES: Final[dict[str, tuple[str, ...]]] = {     "eurollm": (         "utter-project/EuroLLM-1.7B-Instruct",@@ -52,7 +65,10 @@         "Qwen/Qwen3.5-27B",         "Qwen/Qwen3.5-35B-A3B",     ),-    "riva": ("nvidia/Riva-Translate-4B-Instruct",),+    "riva": (+        "nvidia/Riva-Translate-4B-Instruct",+        RIVA_V2_MODEL_NAME,+    ), } SUPPORTED_TRANSLATION_MODELS: Final[tuple[str, ...]] = tuple(     model for series_models in TRANSLATION_MODELS_BY_SERIES.values() for model in series_models@@ -109,11 +125,14 @@ def __init__(         self.split_batch = self.batch_size > 0          self.nmt_model = self.load_model(llm_params)+        self.is_riva_v2 = model_name == RIVA_V2_MODEL_NAME+        self.prefix_tokenizer = self.nmt_model.get_tokenizer() if self.is_riva_v2 else None         self.sampling_params = SamplingParams(**sampling_params)          self.source_language = source_language         self.target_language = target_language         self.prompt_template = self.get_prompt_template(model_name)+        self.use_token_prefix_boundaries = self.is_riva_v2 and self._is_chinese(target_language)         self.waitk = waitk      @staticmethod@@ -131,6 +150,10 @@ def convert_to_dict(params: dict | DictConfig | None) -> dict:             return OmegaConf.to_container(params)         return params +    @staticmethod+    def _is_chinese(language: str) -> bool:+        return language.lower() in CHINESE_LANGUAGE_NAMES+     @staticmethod     def setup_device(device: str, device_id: int) -> tuple[str, int]:         """@@ -189,6 +212,9 @@ def get_prompt_template(model_name: str) -> PromptTemplate:         ):             return QwenReasoningTranslatorPromptTemplate +        if model_name == RIVA_V2_MODEL_NAME:+            return RivaV2TranslatorPromptTemplate+         if model_name in TRANSLATION_MODELS_BY_SERIES["riva"]:             return RivaTranslatorPromptTemplate @@ -204,8 +230,10 @@ def load_model(self, llm_params: dict) -> LLM:             RuntimeError: If model loading fails.         """         try:-            os.environ["CUDA_VISIBLE_DEVICES"] = str(self.device_id)-            local_path = self._get_local_model_path(self.model_name)+            # Respect process-level GPU isolation. Overwriting this after ASR initializes+            # CUDA can move the vLLM worker onto a different physical device.+            os.environ.setdefault("CUDA_VISIBLE_DEVICES", str(self.device_id))+            local_path = self._get_local_model_path(self.model_name, revision=llm_params.get("revision"))             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@@ -217,11 +245,12 @@ def load_model(self, llm_params: dict) -> LLM:         except Exception as e:             raise RuntimeError(f"Model loading failed: {str(e)}") from e -    def _get_local_model_path(self, repo_id):+    def _get_local_model_path(self, repo_id: str, revision: str | None = None):         """         Get local model path from HuggingFace model hub.         Args:             repo_id: (str) repository ID of the model+            revision: (str | None) model revision to resolve in the local cache         Returns:             local_path: (str) local path of the model         Raises:@@ -230,6 +259,7 @@ def _get_local_model_path(self, repo_id):         try:             return snapshot_download(                 repo_id=repo_id,+                revision=revision,                 local_files_only=True,             )         except LocalEntryNotFoundError:@@ -263,18 +293,46 @@ def translate_batch(         for src_lang, tgt_lang, src_prefix, tgt_prefix, src_context, tgt_context in zip(             src_langs, tgt_langs, asr_transcripts, prefixes, src_contexts, tgt_contexts         ):-            text = self.prompt_template.format(src_lang, tgt_lang, src_prefix, tgt_prefix, src_context, tgt_context)+            if self.is_riva_v2:+                messages = self.prompt_template.messages(+                    src_lang, tgt_lang, src_prefix, tgt_prefix, src_context, tgt_context+                )+                has_target_prefix = bool(tgt_prefix.strip())+                text = self.prefix_tokenizer.apply_chat_template(+                    messages,+                    tokenize=False,+                    add_generation_prompt=not has_target_prefix,+                    continue_final_message=has_target_prefix,+                )+            else:+                text = self.prompt_template.format(+                    src_lang, tgt_lang, src_prefix, tgt_prefix, src_context, tgt_context+                )             input_texts.append(text)          outputs = self.nmt_model.generate(input_texts, self.sampling_params, use_tqdm=False)         translations = []-        for tgt_prefix, output in zip(prefixes, outputs):+        for tgt_prefix, tgt_lang, output in zip(prefixes, tgt_langs, outputs):             output_text = output.outputs[0].text-            output_text = self.prompt_template.extract(output_text).strip()+            output_text = self.prompt_template.extract(output_text)             if tgt_prefix:-                translations.append(f"{tgt_prefix} {output_text}")+                if self.is_riva_v2:+                    if self._is_chinese(tgt_lang):+                        separator = ""+                    else:+                        output_text = output_text.strip()+                        separator = ""+                        if (+                            output_text+                            and output_text[0] not in string.punctuation+                            and tgt_prefix[-1] not in "([{'\"«„“"+                        ):+                            separator = " "+                    translations.append(f"{tgt_prefix}{separator}{output_text}")+                else:+                    translations.append(f"{tgt_prefix} {output_text.strip()}")             else:-                translations.append(output_text)+                translations.append(output_text.strip())         return translations      def translate(@@ -314,6 +372,26 @@ def translate(             )         return all_translations +    @staticmethod+    def _token_boundaries(tokenizer, text: str) -> tuple[set[int], list[int]]:+        encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)+        return {end for _, end in encoded["offset_mapping"]}, encoded["input_ids"]++    def _trim_lcp_to_token_boundary(self, previous: str, current: str, lcp: str) -> str:+        """Retreat an LCP to a boundary shared by both model tokenizations."""++        if not lcp:+            return ""+        previous_boundaries, previous_ids = self._token_boundaries(self.prefix_tokenizer, previous)+        current_boundaries, current_ids = self._token_boundaries(self.prefix_tokenizer, current)+        shared_boundaries = previous_boundaries & current_boundaries+        for boundary in sorted((offset for offset in shared_boundaries if offset <= len(lcp)), reverse=True):+            prefix = lcp[:boundary]+            prefix_ids = self.prefix_tokenizer.encode(prefix, add_special_tokens=False)+            if previous_ids[: len(prefix_ids)] == prefix_ids and current_ids[: len(prefix_ids)] == prefix_ids:+                return prefix+        return ""+     def get_prefixes(         self,         asr_transcripts: list[str],@@ -337,8 +415,10 @@ def get_prefixes(             lcp = os.path.commonprefix([prev_trans, trans])             had_leading_space = lcp.startswith(" ") +            if self.use_token_prefix_boundaries:+                lcp = self._trim_lcp_to_token_boundary(prev_trans, trans, lcp)             # If lcp happens mid-word, remove generated ending up to the first full word-            if (len(lcp) > 0) and (lcp[-1] not in f"{string.punctuation} "):+            elif (len(lcp) > 0) and (lcp[-1] not in f"{string.punctuation} "):                 lcp = " ".join(lcp.split()[:-1])              # Remove trailing whitespaces

@@ -153,6 +153,104 @@ def extract(cls, response: str) -> str:         return response.split('\n')[0]  +class RivaV2TranslatorPromptTemplate(PromptTemplate):+    """Message builder for the Riva-Translate-4B-Instruct-v2 chat template."""++    LANGUAGE_CODES = {+        "arabic": "ar",+        "brazilian portuguese": "pt-br",+        "bulgarian": "bg",+        "chinese": "zh",+        "croatian": "hr",+        "czech": "cs",+        "danish": "da",+        "dutch": "nl",+        "english": "en",+        "estonian": "et",+        "european portuguese": "pt-pt",+        "european spanish": "es-es",+        "finnish": "fi",+        "french": "fr",+        "german": "de",+        "greek": "el",+        "hindi": "hi",+        "hungarian": "hu",+        "indonesian": "id",+        "italian": "it",+        "japanese": "ja",+        "korean": "ko",+        "latin american spanish": "es-us",+        "latam spanish": "es-us",+        "latvian": "lv",+        "lithuanian": "lt",+        "norwegian": "no",+        "polish": "pl",+        "portuguese": "pt",+        "romanian": "ro",+        "russian": "ru",+        "simplified chinese": "zh",+        "slovak": "sk",+        "slovenian": "sl",+        "spanish": "es",+        "swedish": "sv",+        "thai": "th",+        "traditional chinese": "zh-tw",+        "turkish": "tr",+        "ukrainian": "uk",+        "vietnamese": "vi",+    }++    @classmethod+    def messages(+        cls,+        src_lang: str,+        tgt_lang: str,+        src_prefix: str,+        tgt_prefix: str,+        src_context: str = "",+        tgt_context: str = "",+    ) -> list[dict[str, str]]:+        """Build messages to be rendered by the model-provided chat template."""++        source_code = cls.LANGUAGE_CODES.get(src_lang.lower(), src_lang.lower())+        target_code = cls.LANGUAGE_CODES.get(tgt_lang.lower(), tgt_lang.lower())+        messages = [{"role": "system", "content": f"{source_code}-{target_code}"}]++        source_context = re.sub(r"\s+", " ", src_context).strip()+        target_context = re.sub(r"\s+", " ", tgt_context).strip()+        if bool(source_context) != bool(target_context):+            raise ValueError("Riva v2 source and target context must either both be set or both be empty")+        if source_context:+            messages.extend(+                [+                    {"role": "user", "content": source_context},+                    {"role": "assistant", "content": target_context},+                ]+            )++        messages.append({"role": "user", "content": re.sub(r"\s+", " ", src_prefix).strip()})+        target_prefix = re.sub(r"\s+", " ", tgt_prefix).strip()+        if target_prefix:+            messages.append({"role": "assistant", "content": target_prefix})+        return messages++    @classmethod+    def format(+        cls,+        src_lang: str,+        tgt_lang: str,+        src_prefix: str,+        tgt_prefix: str,+        src_context: str = "",+        tgt_context: str = "",+    ) -> str:+        raise RuntimeError("Riva v2 prompts must be rendered with the model tokenizer")++    @classmethod+    def extract(cls, response: str) -> str:+        return response.split('\n')[0]++ class QwenReasoningTranslatorPromptTemplate(PromptTemplate):     """     Chat-style prompt template for Qwen Reasoning model to perform translation.

@@ -0,0 +1,191 @@+# 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.++from types import SimpleNamespace+from unittest.mock import patch++import pytest++pytest.importorskip("vllm", reason="Riva Translate v2 tests require vLLM")++from nemo.collections.asr.inference.nmt import llm_translator+from nemo.collections.asr.inference.nmt.llm_translator import RIVA_V2_MODEL_NAME, LLMTranslator+from nemo.collections.asr.inference.nmt.prompts import RivaV2TranslatorPromptTemplate+++class _PairTokenizer:+    def __call__(self, text, add_special_tokens=False, return_offsets_mapping=False):+        offsets = [(start, min(start + 2, len(text))) for start in range(0, len(text), 2)]+        result = {"input_ids": [hash(text[start:end]) for start, end in offsets]}+        if return_offsets_mapping:+            result["offset_mapping"] = offsets+        return result++    def encode(self, text, add_special_tokens=False):+        return self(text, add_special_tokens=add_special_tokens)["input_ids"]+++def test_riva_v2_is_a_supported_translation_model():+    assert LLMTranslator.get_prompt_template(RIVA_V2_MODEL_NAME) is RivaV2TranslatorPromptTemplate+++def test_local_model_resolution_honors_revision():+    calls = []++    def _snapshot_download(**kwargs):+        calls.append(kwargs)+        return "/cached/revision"++    translator = LLMTranslator.__new__(LLMTranslator)++    with patch.object(llm_translator, "snapshot_download", _snapshot_download):+        path = translator._get_local_model_path("nvidia/model", revision="abc123")++    assert path == "/cached/revision"+    assert calls == [{"repo_id": "nvidia/model", "revision": "abc123", "local_files_only": True}]+++def test_model_loading_preserves_process_level_gpu_isolation():+    loaded = []++    translator = LLMTranslator.__new__(LLMTranslator)+    translator.model_name = "nvidia/model"+    translator.device_id = 0++    with (+        patch.dict(llm_translator.os.environ, {"CUDA_VISIBLE_DEVICES": "1"}),+        patch.object(llm_translator, "LLM", lambda **kwargs: loaded.append(kwargs) or object()),+        patch.object(translator, "_get_local_model_path", lambda *args, **kwargs: None),+    ):+        translator.load_model({})++        assert llm_translator.os.environ["CUDA_VISIBLE_DEVICES"] == "1"+    assert loaded == [{"model": "nvidia/model"}]+++def test_riva_v2_messages_use_language_pair_and_aligned_context():+    messages = RivaV2TranslatorPromptTemplate.messages(+        "English", "German", "Current source", "Aktuelle", "Earlier source", "Fruehere Ausgabe"+    )++    assert messages == [+        {"role": "system", "content": "en-de"},+        {"role": "user", "content": "Earlier source"},+        {"role": "assistant", "content": "Fruehere Ausgabe"},+        {"role": "user", "content": "Current source"},+        {"role": "assistant", "content": "Aktuelle"},+    ]+++def test_riva_v2_uses_model_chat_template_for_prefix_continuation():+    calls = []++    class _Tokenizer:+        def apply_chat_template(self, messages, **kwargs):+            calls.append((messages, kwargs))+            return "tokenizer-rendered-prompt"++    class _Model:+        def generate(self, prompts, sampling_params, use_tqdm):+            assert prompts == ["tokenizer-rendered-prompt"]+            return [SimpleNamespace(outputs=[SimpleNamespace(text=".")])]++    translator = LLMTranslator.__new__(LLMTranslator)+    translator.is_riva_v2 = True+    translator.prompt_template = RivaV2TranslatorPromptTemplate+    translator.prefix_tokenizer = _Tokenizer()+    translator.nmt_model = _Model()+    translator.sampling_params = object()+    translator.use_token_prefix_boundaries = False++    result = translator.translate_batch(["Current source"], ["Die Uebersetzung"], ["en"], ["de"], [""], [""])++    assert result == ["Die Uebersetzung."]+    assert calls[0][1] == {+        "tokenize": False,+        "add_generation_prompt": False,+        "continue_final_message": True,+    }+++def test_riva_v2_opens_assistant_turn_without_target_prefix():+    calls = []++    class _Tokenizer:+        def apply_chat_template(self, messages, **kwargs):+            calls.append((messages, kwargs))+            return "tokenizer-rendered-prompt"++    class _Model:+        def generate(self, prompts, sampling_params, use_tqdm):+            return [SimpleNamespace(outputs=[SimpleNamespace(text="Uebersetzung")])]++    translator = LLMTranslator.__new__(LLMTranslator)+    translator.is_riva_v2 = True+    translator.prompt_template = RivaV2TranslatorPromptTemplate+    translator.prefix_tokenizer = _Tokenizer()+    translator.nmt_model = _Model()+    translator.sampling_params = object()+    translator.use_token_prefix_boundaries = False++    result = translator.translate_batch(["Current source"], [""], ["en"], ["de"], [""], [""])++    assert result == ["Uebersetzung"]+    assert calls[0][1] == {+        "tokenize": False,+        "add_generation_prompt": True,+        "continue_final_message": False,+    }+++def test_riva_v2_joins_prefixes_using_each_request_target_language():+    class _Tokenizer:+        def apply_chat_template(self, messages, **kwargs):+            return "tokenizer-rendered-prompt"++    class _Model:+        def generate(self, prompts, sampling_params, use_tqdm):+            return [+                SimpleNamespace(outputs=[SimpleNamespace(text="开始")]),+                SimpleNamespace(outputs=[SimpleNamespace(text="Konferenz")]),+            ]++    translator = LLMTranslator.__new__(LLMTranslator)+    translator.is_riva_v2 = True+    translator.prompt_template = RivaV2TranslatorPromptTemplate+    translator.prefix_tokenizer = _Tokenizer()+    translator.nmt_model = _Model()+    translator.sampling_params = object()++    result = translator.translate_batch(+        ["Current source", "Current source"],+        ["会议", "Die"],+        ["en", "en"],+        ["zh", "de"],+        ["", ""],+        ["", ""],+    )++    assert result == ["会议开始", "Die Konferenz"]+++def test_riva_v2_token_lcp_retreats_to_shared_model_token_boundary():+    translator = LLMTranslator.__new__(LLMTranslator)+    translator.prefix_tokenizer = _PairTokenizer()+    translator.use_token_prefix_boundaries = True+    translator.waitk = -1++    prefix = translator.get_prefixes(["source"], ["会议开幕"], ["会议开始"])[0]++    assert prefix == "会议"