NVIDIA-NeMo/Speech · #15940

Add inference support for the Unified multilingual RNNT model with lang prompt

andrusenkoau · merged Oct 23, 202618 files · 1246 + / 221 −
docs/source/asr/api.rst1 + / 1 −
@@ -27,7 +27,7 @@ Model Classes  .. autoclass:: nemo.collections.asr.models.EncDecRNNTBPEModelWithPrompt     :show-inheritance:-    :members: transcribe, set_inference_prompt, initialize_prompt_feature, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact+    :members: transcribe, forward, resolve_prompt_language, set_inference_prompt, initialize_prompt_feature, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact   .. autoclass:: nemo.collections.asr.models.EncDecHybridRNNTCTCBPEModelWithPrompt
docs/source/asr/configs.rst28 + / 8 −
@@ -863,28 +863,43 @@ The model introduces several prompt-specific configuration parameters in the ``m  **Dataset Configuration:** -The model requires training data with prompt annotations when using Lhotse datasets:+The model requires training data with language annotations when using Lhotse datasets:  .. code-block:: yaml    model:     train_ds:       use_lhotse: true       initialize_prompt_feature: true-      prompt_field: "target_lang"     # Field name for prompt extraction+      lang_field: target_lang         # Manifest key holding the utterance language       prompt_dictionary: ${model.model_defaults.prompt_dictionary}       num_prompts: ${model.model_defaults.num_prompts}-      +      default_prompt_mode: unified    # langID | auto | unified+      unified_auto_ratio: 0.5         # Probability of the "auto" prompt in unified mode+     validation_ds:       use_lhotse: true       initialize_prompt_feature: true-      prompt_field: "target_lang"+      lang_field: target_lang       prompt_dictionary: ${model.model_defaults.prompt_dictionary}       num_prompts: ${model.model_defaults.num_prompts}+      default_prompt_mode: langID++The prompt index for each utterance is chosen by its prompt mode:++- ``langID``: the utterance's language.+- ``auto``: the language-agnostic ``auto`` prompt.+- ``unified``: ``auto`` with probability ``unified_auto_ratio``, otherwise the utterance's language.+  This is the default; it trains one model for both prompted and language-agnostic decoding.++A Lhotse input config can set the mode per dataset with ``tags: {prompt_mode: langID}``; otherwise+``default_prompt_mode`` applies. Transcription always uses ``langID``.  **Manifest Format:** -Training manifests should include prompt information:+The language is read from the manifest key named by ``lang_field`` (``lang`` by default; the example+configs use ``target_lang``). It must be a key of ``prompt_dictionary``, or a bare code such as ``en``+that matches exactly one locale key; unsupported languages raise an error.  .. code-block:: json @@ -941,7 +956,7 @@ variant:         'fr-FR': 2,         'es-ES': 3,         # ... additional language codes ...-        'auto': 127,                   # Per-sample dynamic language (read from manifest)+        'auto': 127,                   # Language-agnostic prompt       }  **Dataset Configuration:**@@ -955,16 +970,21 @@ The model uses the same index-based Lhotse dataset     train_ds:       use_lhotse: true       initialize_prompt_feature: true-      prompt_field: "target_lang"     # Field name for per-sample prompt extraction+      lang_field: target_lang       prompt_dictionary: ${model.model_defaults.prompt_dictionary}       num_prompts: ${model.model_defaults.num_prompts}+      default_prompt_mode: unified    # langID | auto | unified+      unified_auto_ratio: 0.5      validation_ds:       use_lhotse: true       initialize_prompt_feature: true-      prompt_field: "target_lang"+      lang_field: target_lang       prompt_dictionary: ${model.model_defaults.prompt_dictionary}       num_prompts: ${model.model_defaults.num_prompts}+      default_prompt_mode: langID++Prompt modes behave as described for the hybrid model.  **Manifest Format:** 
docs/source/asr/inference.rst7 + / 2 −
@@ -275,8 +275,13 @@ flag directly:         decoder_type=rnnt \         strip_lang_tags=true -Use ``target_lang=<lang-code>`` to pin every sample in the batch to one language, or-``target_lang=auto`` to read the per-sample ``target_lang`` field from each manifest entry.+``target_lang`` must be a key of the model's ``prompt_dictionary``. A bare language code such as+``en`` is also accepted when it matches exactly one locale key (for example ``en-US``). An+unsupported language raises an error. ``target_lang=auto`` (the default when ``target_lang`` is+unset) selects the language-agnostic prompt; quality on languages the model was not trained on is+not guaranteed. This script applies one prompt to every sample. For per-utterance languages, use+``model.transcribe()``, which reads each manifest entry's ``lang`` field when ``target_lang`` is+unset and uses ``auto`` for entries without one. Setting ``strip_lang_tags=true`` removes ``<xx-XX>`` language tags from the decoded text (pattern is customizable via ``lang_tag_pattern``). 
docs/source/asr/results.rst9 + / 3 −
@@ -317,9 +317,14 @@ The Hybrid-Transducer-CTC model with prompt conditioning (``EncDecHybridRNNTCTCB   # Transcribe with language prompts   transcriptions = model.transcribe(       paths2audio_files=["audio1.wav", "audio2.wav"],-      target_lang="en-US"  # Specify target language+      target_lang="en-US"  # Language of the audio; must be supported by the model   ) +``target_lang`` must be a language from the model's ``prompt_dictionary``; an unsupported language+raises an error. If the language is unknown, pass ``target_lang="auto"`` or leave it unset (``auto``+is the default). The model is also trained with this language-agnostic prompt, so it still works,+but accuracy is usually slightly lower than with the correct language.+ **Configuration and Training:**  - Training script: ``<NeMo_git_root>/examples/asr/asr_hybrid_transducer_ctc/speech_to_text_hybrid_rnnt_ctc_bpe_prompt.py``@@ -338,8 +343,9 @@ auxiliary CTC head.  - Cache-aware streaming RNN-T with language-ID prompt conditioning - Concatenation-based feature fusion with learned projection (shared with the hybrid variant)-- Per-sample dynamic language selection from the manifest (``target_lang=auto``) or-  fixed-language batches (e.g. ``target_lang=en-US``)+- Fixed-language decoding (e.g. ``target_lang=en-US``), the language-agnostic prompt+  (``target_lang=auto``), or, when ``target_lang`` is unset, per-utterance languages read from the+  manifest ``lang`` field. Languages missing from the model's ``prompt_dictionary`` raise an error. - Optional automatic stripping of language tags from decoded text via ``strip_lang_tags``  **Usage Example:**
examples/asr/asr_chunked_inference/rnnt/speech_to_text_streaming_infer_rnnt.py62 + / 3 −
@@ -21,6 +21,17 @@ This is especially useful for models such as Conformers, which have quadratic time and memory scaling with audio duration. +Supported models:+- Standard (offline) RNNT / Hybrid RNNT-CTC models: transcribed here in a buffered/streaming fashion.+- Prompt-conditioned multilingual RNNT models (`EncDecRNNTBPEModelWithPrompt`): additionally inject a+  language-ID prompt into the encoder output. The prompt encodes the language spoken in the audio,+  selected per utterance from each manifest record's `lang_field` (default "lang", the same field+  offline transcription reads, e.g. "en" or "en-US"). Passing `target_lang` overrides that field+  for every utterance (handy to force one language). A record without a language uses the+  language-agnostic "auto" prompt; a language the model does not support raises an error.+  Both are ignored by models without prompt conditioning, and are separate from `langid`, which only+  selects the number-to-words locale used when `clean_groundtruth_text=True`.+ The difference between streaming and buffered inference is the chunk size (or the latency of inference). Buffered inference will use large chunk sizes (5-10 seconds) + some additional right for context. Streaming inference will use small chunk sizes (0.1 to 0.25 seconds) + some additional right buffer for context.@@ -31,7 +42,7 @@ - long file transcription: in most cases 10-10-5 (10s left, 10s chunk, 5s right) will give results similar to offline - streaming with 4s latency: 10-2-2 is usually similar or better than 10-0.16-3.84 and significantly faster -Example usage:+Example usage (standard model):  ```shell python speech_to_text_streaming_infer_rnnt.py \@@ -47,6 +58,9 @@     clean_groundtruth_text=False \     langid='en' ```++For prompt-conditioned models, add `target_lang='en'` to force one language for every record,+instead of reading each record's `lang_field`. """ import copy import glob@@ -143,6 +157,15 @@ class TranscriptionConfig:         True  # whether to use the att_context_size as chunk size (important for extra-low latency)     ) +    # Language-ID prompt for prompt-conditioned models: the language spoken in the audio (e.g. "en",+    # "de-DE", or the language-agnostic "auto"), applied to every utterance. By default (None) each+    # record uses its own `lang_field`, and "auto" when it names none. A language the model does not+    # support raises an error. Ignored by models without prompt support.+    target_lang: Optional[str] = None+    # "lang" matches the Lhotse manifest field that offline transcription reads, so the same manifest+    # conditions both paths identically.+    lang_field: str = "lang"+     # Set `cuda` to int to define CUDA device. If 'None', will look for CUDA     # device anyway, and do inference on CPU only if CUDA device is not found.     # If `cuda` is a negative number, inference will be on CPU only.@@ -394,10 +417,36 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:         in_order=True,     ) +    # Language-ID prompt setup. The prompt encodes the spoken language, resolved per utterance: an+    # explicit `target_lang` overrides every record (e.g. to force a language); otherwise each+    # record's `lang_field` is used, and "auto" when the record names none. A language the model does+    # not support raises before decoding starts. Indices follow `records` order (matching the+    # shuffle=False, in_order=True dataloader), so a running offset maps each batch to its languages.+    # The model builds the one-hot from these indices at the encoder's own length, once per chunk.+    prompt_indices = None+    if getattr(asr_model, 'concat', False):+        per_record_langs = [+            cfg.target_lang if cfg.target_lang is not None else record.get(cfg.lang_field, None) for record in records+        ]+        resolved = {lang: asr_model.resolve_prompt_language(lang) for lang in set(per_record_langs)}+        prompt_indices = torch.tensor(+            [resolved[lang] for lang in per_record_langs], dtype=torch.long, device=map_location+        )+        source = (+            f"target_lang='{cfg.target_lang}'" if cfg.target_lang is not None else f"lang_field='{cfg.lang_field}'"+        )+        logging.info(+            f"Language-ID prompt conditioning enabled for {len(per_record_langs)} utterances from {source}; "+            f"resolved languages -> indices: {resolved}"+        )+     timer = SimpleTimer()     with torch.no_grad(), torch.inference_mode():         all_hyps = []         audio_data: AudioBatch+        # Offset into `prompt_indices`; advanced by each batch since the dataloader preserves+        # `records` order (shuffle=False, in_order=True, drop_last=False).+        prompt_offset = 0         timer.start(device=map_location)         for audio_data in tqdm(audio_dataloader):             # get audio@@ -407,6 +456,12 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:             batch_size = audio_batch.shape[0]             device = audio_batch.device +            # This batch's language-ID indices (one row per utterance, constant across chunks).+            batch_prompt_indices = None+            if prompt_indices is not None:+                batch_prompt_indices = prompt_indices[prompt_offset : prompt_offset + batch_size]+                prompt_offset += batch_size+             # add biasing requests to the decoder             if use_per_stream_biasing:                 multi_biasing_ids = torch.full([batch_size], fill_value=-1, dtype=torch.long, device=map_location)@@ -460,12 +515,16 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:                     is_last_chunk=is_last_chunk,                     is_last_chunk_batch=is_last_chunk_batch,                 )-                 # get encoder output using full buffer [left-chunk-right]-                encoder_output, encoder_output_len = asr_model(+                forward_kwargs = dict(                     input_signal=buffer.samples,                     input_signal_length=buffer.context_size_batch.total(),                 )+                if batch_prompt_indices is not None:+                    forward_kwargs["language_prompt_indices"] = batch_prompt_indices++                encoder_output, encoder_output_len = asr_model(**forward_kwargs)+                 encoder_output = encoder_output.transpose(1, 2)  # [B, T, C]                 # remove extra context from encoder_output (leave only frames corresponding to the chunk)                 encoder_context = buffer.context_size.subsample(factor=encoder_frame2audio_samples)
examples/asr/transcribe_speech.py11 + / 0 −
@@ -189,6 +189,13 @@ class TranscriptionConfig:     # att_context_size can be set for cache-aware streaming models with multiple look-aheads     att_context_size: Optional[list] = None +    # Language-ID prompt for prompt-conditioned ASR RNNT models: the language spoken in the audio+    # (e.g. "en", "de-DE", or the language-agnostic "auto"), applied to every utterance. By default+    # (null) each utterance uses its own manifest language, read from `gt_lang_attr_name`, and "auto"+    # when it names none. A language the model does not support raises an error. Ignored by models+    # without prompt conditioning.+    target_lang: Optional[str] = None+     # Use this for model-specific changes before transcription     model_change: ModelChangeConfig = field(default_factory=ModelChangeConfig) @@ -420,6 +427,10 @@ def main(cfg: TranscriptionConfig) -> Union[TranscriptionConfig, List[Hypothesis             override_cfg.timestamps = cfg.timestamps             if hasattr(override_cfg, "prompt"):                 override_cfg.prompt = parse_multitask_prompt(OmegaConf.to_container(cfg.prompt))+            # Prompt-conditioned RNNT models: forward the language-ID prompt. None is forwarded too,+            # since it asks for each utterance's own manifest language.+            if hasattr(override_cfg, "target_lang"):+                override_cfg.target_lang = cfg.target_lang              device = next(asr_model.parameters()).device             for run_step in range(cfg.warmup_steps + cfg.run_steps):
nemo/collections/asr/data/audio_to_text_lhotse_prompt_index.py33 + / 24 −
@@ -26,6 +26,7 @@ from lhotse.dataset import AudioSamples from lhotse.dataset.collation import collate_vectors +from nemo.collections.asr.parts.utils.prompt_language_utils import AUTO_PROMPT_LANGUAGE, resolve_prompt_index from nemo.collections.common.data.lhotse.audio_loading import LhotseAudioLoadingDatasetMixin from nemo.collections.common.tokenizers.aggregate_tokenizer import TokenizerWrapper from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec@@ -46,7 +47,7 @@ class LhotseSpeechToTextBpeDatasetWithPromptIndex(LhotseAudioLoadingDatasetMixin         audio_signal_length: Audio lengths [B]         transcripts: Token IDs [B, T]         transcript_length: Token lengths [B]-        prompt_indices: Language ID indices [B] (NOT full tensors)+        language_prompt_indices: Language ID indices [B] (NOT full tensors)     """      @property@@ -56,7 +57,7 @@ def output_types(self) -> Optional[Dict[str, NeuralType]]:             'audio_signal_length': NeuralType(tuple('B'), LengthsType()),             'transcripts': NeuralType(('B', 'T'), LabelsType()),             'transcript_length': NeuralType(tuple('B'), LengthsType()),-            'prompt_indices': NeuralType(tuple('B'), LabelsType()),  # Just indices, not full tensors+            'language_prompt_indices': NeuralType(tuple('B'), LabelsType()),  # Just indices, not full tensors         }      def __init__(self, tokenizer: TokenizerSpec, cfg: Dict) -> None:@@ -72,33 +73,41 @@ def __init__(self, tokenizer: TokenizerSpec, cfg: Dict) -> None:          self.num_prompts = cfg.get('num_prompts', 128) -        # Per-dataset prompt mode is read from cut.custom["prompt_mode"] at runtime.-        # Supported values:-        #   "langID"  — always pass the real language ID-        #   "auto"    — always pass auto (101)-        #   "unified" — randomize: auto with probability unified_auto_ratio, else lang ID+        # Per-dataset prompt mode is read from cut.custom["prompt_mode"] at runtime, falling back to+        # `default_prompt_mode`. Supported values:+        #   "langID"  — always pass the cut's language ID. Transcription always uses this mode.+        #   "auto"    — always pass the language-agnostic "auto" prompt.+        #   "unified" — training default: "auto" with probability unified_auto_ratio, else the+        #               language ID, so that one model learns both prompted and agnostic decoding.         # Set via lhotse input_cfg tags, e.g.  tags: { prompt_mode: langID }         self.prompt_mode_field = cfg.get('prompt_mode_field', 'prompt_mode')         self.default_prompt_mode = cfg.get('default_prompt_mode', 'unified')         self.unified_auto_ratio = cfg.get('unified_auto_ratio', 0.5)          # Index used for the language-agnostic / auto prompt-        self.auto_index = self.prompt_dict.get('auto', 101)+        self.auto_index = self.prompt_dict.get(AUTO_PROMPT_LANGUAGE, 101)++        # Language that replaces every cut's own language (set by transcription for a run-wide+        # `target_lang`), and language for cuts that name none. Training leaves both unset, so a cut+        # without a language is an error there.+        self.target_lang = cfg.get('target_lang', None)+        self.default_lang = cfg.get('default_lang', None)          logging.info(             f"LhotseSpeechToTextBpeDatasetWithPromptIndex: "             f"default_prompt_mode={self.default_prompt_mode}, "-            f"unified_auto_ratio={self.unified_auto_ratio}"+            f"unified_auto_ratio={self.unified_auto_ratio}, "+            f"target_lang={self.target_lang}, default_lang={self.default_lang}"         ) -    def _get_prompt_index(self, prompt_key: str) -> int:-        """Maps prompt keys to indices using the prompt dictionary."""-        if prompt_key not in self.prompt_dict:-            available_keys = list(self.prompt_dict.keys())-            raise ValueError(-                f"Unknown prompt key: '{prompt_key}'. Available: {available_keys[:10]}{'...' if len(available_keys) > 10 else ''}"-            )-        return self.prompt_dict[prompt_key]+    def _get_language_index(self, cut) -> int:+        """Prompt index for the cut's language; raises if the model does not support it."""+        language = self.target_lang if self.target_lang is not None else cut.supervisions[0].language+        if language is None:+            if self.default_lang is None:+                raise ValueError(f"Cut '{cut.id}' has no language, and no default_lang is configured.")+            language = self.default_lang+        return resolve_prompt_index(language, self.prompt_dict)      def _get_prompt_mode(self, cut) -> str:         """Resolve the prompt_mode for a cut from its custom tags."""@@ -114,34 +123,34 @@ def _get_prompt_index_for_cut(self, cut) -> int:          Behaviour depends on prompt_mode (set per-dataset via lhotse         input_cfg tags, falling back to ``default_prompt_mode``):-            "langID"  — always return the real language ID-                        (use for inference / language-forced tasks)+            "langID"  — always return the cut's language ID+                        (``target_lang`` replaces it when set; ``default_lang`` covers cuts without one)             "auto"    — always return auto index (language-agnostic)             "unified" — return auto with probability unified_auto_ratio,                         otherwise the real language ID         """         mode = self._get_prompt_mode(cut)          if mode == 'langID':-            return self._get_prompt_index(cut.supervisions[0].language)+            return self._get_language_index(cut)         elif mode == 'auto':             return self.auto_index         elif mode == 'unified':             if random.random() < self.unified_auto_ratio:                 return self.auto_index-            return self._get_prompt_index(cut.supervisions[0].language)+            return self._get_language_index(cut)         else:             logging.warning(f"Unknown prompt_mode '{mode}', falling back to unified")             if random.random() < self.unified_auto_ratio:                 return self.auto_index-            return self._get_prompt_index(cut.supervisions[0].language)+            return self._get_language_index(cut)      def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:         audio, audio_lens, cuts = self.load_audio_with_cuts(cuts)         tokens = [torch.as_tensor(self.tokenizer(c.supervisions[0].text, c.supervisions[0].language)) for c in cuts]          # Get prompt indices (just the language ID per sample, NOT full tensors)-        prompt_indices = torch.tensor([self._get_prompt_index_for_cut(c) for c in cuts], dtype=torch.long)+        language_prompt_indices = torch.tensor([self._get_prompt_index_for_cut(c) for c in cuts], dtype=torch.long)          # Create final tensors         token_lens = torch.tensor([t.size(0) for t in tokens], dtype=torch.long)@@ -152,5 +161,5 @@ def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:             audio_lens,  # Audio lengths [B]             tokens,  # Text tokens [B, T]             token_lens,  # Token lengths [B]-            prompt_indices,  # Language ID indices [B] - model creates full tensor+            language_prompt_indices,  # Language ID indices [B] - model creates full tensor         )
nemo/collections/asr/inference/model_wrappers/rnnt_inference_wrapper.py17 + / 12 −
@@ -73,15 +73,19 @@ def get_subsampling_factor(self) -> int:         return self.asr_model.encoder.subsampling_factor      def encode(-        self, processed_signal: Tensor, processed_signal_length: Tensor, prompt_vectors: Tensor | None = None+        self,+        processed_signal: Tensor,+        processed_signal_length: Tensor,+        language_prompt_indices: Tensor | None = None,     ) -> tuple[Tensor, Tensor]:         """         Get encoder output from the model. It is used for streaming inference.         Args:             processed_signal: (Tensor) processed signal. Shape is torch.Size([B, C, T]).             processed_signal_length: (Tensor) processed signal length. Shape is torch.Size([B]).-            prompt_vectors: (Tensor | None) Optional prompt vectors for multilingual models.-                Shape can be torch.Size([B, num_prompts]) or torch.Size([B, T_enc, num_prompts]) if already expanded.+            language_prompt_indices: (Tensor | None) Optional per-sample language-ID indices for+                multilingual models. Shape is torch.Size([B]). The model expands them across the+                time axis using its own encoder output length.         Returns:             (tuple[Tensor, Tensor]) encoder output and encoder output length of shape torch.Size([B, T, D]), torch.Size([B]).         """@@ -102,8 +106,8 @@ def encode(                 'processed_signal': processed_signal.to(self.cast_dtype),                 'processed_signal_length': processed_signal_length,             }-            if prompt_vectors is not None:-                model_args['prompt'] = prompt_vectors+            if language_prompt_indices is not None:+                model_args['language_prompt_indices'] = language_prompt_indices.to(self.device)              forward_outs = self.asr_model(**model_args) @@ -126,23 +130,24 @@ def decode(self, encoded: Tensor, encoded_len: Tensor, partial_hypotheses: list)         return best_hyp      def encode_with_prompts(-        self, processed_signal: Tensor, processed_signal_length: Tensor, prompt_vectors: Tensor+        self, processed_signal: Tensor, processed_signal_length: Tensor, language_prompt_indices: Tensor     ) -> tuple[Tensor, Tensor]:         """         Convenience wrapper for prompt-enabled encoding.-        Expands prompt vectors across the time dimension before calling encode.++        The indices are expanded across the time dimension by the model, after encoding. Dividing+        the feature length by the subsampling factor here would instead disagree with the encoder+        whenever that length is not an exact multiple of it: subsampling rounds each stage up, so+        the estimate can come out one frame short and the prompt then fails to concatenate.         Args:             processed_signal: (Tensor) processed signal. Shape is torch.Size([B, C, T]).             processed_signal_length: (Tensor) processed signal length. Shape is torch.Size([B]).-            prompt_vectors: (Tensor) prompt vectors. Shape is torch.Size([B, num_prompts]).+            language_prompt_indices: (Tensor) per-sample language-ID indices. Shape is torch.Size([B]).         Returns:             (tuple[Tensor, Tensor]) encoder output and encoder output length.         """-        encoder_time_steps = processed_signal.shape[2] // self.get_subsampling_factor()-        # Expand prompts: [B, num_prompts] -> [B, T_enc, num_prompts]-        prompt_vectors = prompt_vectors.unsqueeze(1).expand(-1, encoder_time_steps, -1)         return self.encode(             processed_signal=processed_signal,             processed_signal_length=processed_signal_length,-            prompt_vectors=prompt_vectors,+            language_prompt_indices=language_prompt_indices,         )
nemo/collections/asr/inference/pipelines/base_pipeline.py23 + / 30 −
@@ -46,6 +46,7 @@ ) from nemo.collections.asr.inference.utils.progressbar import ProgressBar from nemo.collections.asr.inference.utils.text_segment import TextSegment+from nemo.collections.asr.parts.utils.prompt_language_utils import resolve_prompt_index from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec  if TYPE_CHECKING:@@ -521,41 +522,21 @@ def _load_prompt_config(self) -> dict:          return {} -    def _resolve_prompt_index(self, language_code: str) -> int:+    def _resolve_prompt_index(self, language_code: str | None) -> int:         """-        Resolve language_code to a strict prompt index; raise if invalid.+        Resolve language_code to a prompt index, with the rules shared by the prompt-conditioned models.         Args:-            language_code: (str) Language code to resolve (e.g., "en-US", "es-ES").+            language_code: (str | None) Language code to resolve (e.g., "en", "es-ES"), or None for+                the language-agnostic "auto" prompt.         Returns:             (int) Prompt index corresponding to the language code.         Raises:             RuntimeError: If prompt configuration is missing.-            ValueError: If language_code is not found in prompt dictionary.+            ValueError: If the language is not supported by the model, or is ambiguous.         """         if not hasattr(self, '_prompt_config') or not self._prompt_config:             raise RuntimeError("Prompt configuration is missing for a prompt-enabled model.")-        prompt_dict = self._prompt_config['prompt_dict']-        lang_index = prompt_dict.get(language_code, None)-        if lang_index is None:-            raise ValueError(-                f"Language code '{language_code}' not found in prompt dictionary. "-                f"Available languages: {list(prompt_dict.keys())}"-            )-        return lang_index--    def _resolve_default_language_code(self) -> str:-        """-        Pick the language used when a request does not specify one.--        Prefers the model's automatic language-detection prompt when available, so that a-        multilingual model does not silently transcribe every language as English.-        Returns:-            (str) "auto" when the model's prompt dictionary supports it, otherwise "en-US".-        """-        if not getattr(self, '_prompt_config', None):-            return "en-US"-        prompt_dict = self._prompt_config['prompt_dict']-        return "auto" if "auto" in prompt_dict else "en-US"+        return resolve_prompt_index(language_code or None, self._prompt_config['prompt_dict'])      def _create_one_hot_prompts(self, indices: Tensor) -> Tensor:         """@@ -568,21 +549,33 @@ def _create_one_hot_prompts(self, indices: Tensor) -> Tensor:         num_prompts = self._prompt_config['num_prompts']         return torch.nn.functional.one_hot(indices, num_classes=num_prompts).to(self._prompt_config['compute_dtype']) -    def _build_prompt_vectors(self, states: list) -> Tensor:+    def _build_prompt_indices(self, states: list) -> Tensor:         """-        Build prompt vectors for a batch of states using one-hot encoding.+        Collect the per-stream prompt indices for a batch of states.         Args:             states: (list) List of streaming states.         Returns:-            (Tensor) Prompt vectors of shape [B, num_prompts].+            (Tensor) Prompt indices of shape [B].         Raises:             ValueError: If any prompt index is out of range.         """         indices = torch.tensor([getattr(s, 'prompt_idx', 0) for s in states], device=self.device, dtype=torch.long)         num_prompts = self._prompt_config['num_prompts']         if torch.any((indices < 0) | (indices >= num_prompts)):             raise ValueError("Found out-of-range prompt index in batch.")-        return self._create_one_hot_prompts(indices)+        return indices++    def _build_prompt_vectors(self, states: list) -> Tensor:+        """+        Build prompt vectors for a batch of states using one-hot encoding.+        Args:+            states: (list) List of streaming states.+        Returns:+            (Tensor) Prompt vectors of shape [B, num_prompts].+        Raises:+            ValueError: If any prompt index is out of range.+        """+        return self._create_one_hot_prompts(self._build_prompt_indices(states))      def run(         self,
nemo/collections/asr/inference/pipelines/buffered_rnnt_pipeline.py15 + / 19 −
@@ -132,6 +132,8 @@ def init_parameters(self, cfg: DictConfig) -> None:         self.residue_tokens_at_end = cfg.endpointing.residue_tokens_at_end         self.word_boundary_tolerance = cfg.streaming.word_boundary_tolerance         self.return_tail_result = cfg.return_tail_result+        # Language used to select the prompt for prompt-conditioned models; None selects "auto".+        self.default_language_code = cfg.get('lang', None)         self.tokens_to_move = self.punctuation_ids.union(self.language_token_ids)          self.zero_encoded = self.init_zero_enc() if self.right_padding else None@@ -200,16 +202,15 @@ def init_zero_enc(self) -> Tensor:         )          if self.prompt_enabled:-            # Use "en-US" as the default prompt for zero encoding-            # This region is sliced out before decoding, so language choice doesn't matter-            default_prompt_idx = self._resolve_prompt_index("en-US")+            # This region is sliced out before decoding, so the language choice doesn't matter; any+            # index the model defines will do.+            default_prompt_idx = next(iter(self._prompt_config['prompt_dict'].values()))             prompt_indices = torch.tensor([default_prompt_idx], device=self.device, dtype=torch.long)-            prompt_vector = self._create_one_hot_prompts(prompt_indices)  # [1, num_prompts]              zero_encoded, _ = self.asr_model.encode_with_prompts(                 processed_signal=zero_features,                 processed_signal_length=zero_features_len,-                prompt_vectors=prompt_vector,+                language_prompt_indices=prompt_indices,             )         else:             zero_encoded, _ = self.asr_model.encode(@@ -235,16 +236,13 @@ def create_state(self, options: ASRRequestOptions) -> RNNTStreamingState:             default_target_language=self.nmt_model.target_language if self.nmt_enabled else None,             default_stop_history_eou=self.stop_history_eou_in_milliseconds,             default_asr_output_granularity=self.asr_output_granularity,-            default_language_code="en-US" if self.prompt_enabled else None,+            default_language_code=self.default_language_code if self.prompt_enabled else None,         )         state.set_options(new_options) -        # Create per-stream prompt index for prompt-enabled models+        # Create per-stream prompt index for prompt-enabled models; no language selects "auto".         if self.prompt_enabled:-            lang_code = getattr(new_options, "language_code", None)-            if not isinstance(lang_code, str) or len(lang_code) == 0:-                raise ValueError("Prompt-enabled model requires a valid language_code in request options.")-            prompt_idx = self._resolve_prompt_index(lang_code)+            prompt_idx = self._resolve_prompt_index(getattr(new_options, "language_code", None))             state.set_prompt_index(prompt_idx)          return state@@ -330,16 +328,15 @@ def encode_raw_signals(             expected_feature_buffer_len=self.expected_feature_buffer_len,         ) -        # Build prompt vectors if prompts are enabled+        # Build prompt indices if prompts are enabled         if self.prompt_enabled:             requests_states = [self.get_state(f.stream_id) for f in frames]-            prompt_vectors = self._build_prompt_vectors(requests_states)+            prompt_indices = self._build_prompt_indices(requests_states) -            # Use encode_with_prompts which handles dimension expansion             encoded, encoded_len = self.asr_model.encode_with_prompts(                 processed_signal=feature_buffers,                 processed_signal_length=feature_buffer_lens,-                prompt_vectors=prompt_vectors,+                language_prompt_indices=prompt_indices,             )         else:             encoded, encoded_len = self.asr_model.encode(@@ -378,16 +375,15 @@ def encode_processed_signals(         processed_signals = normalize_features(processed_signals, processed_signal_lengths)         processed_signal_lengths = processed_signal_lengths.clamp(max=processed_signals.shape[2]) -        # Build prompt vectors if prompts are enabled+        # Build prompt indices if prompts are enabled         if self.prompt_enabled:             requests_states = [self.get_state(f.stream_id) for f in fbuffers]-            prompt_vectors = self._build_prompt_vectors(requests_states)+            prompt_indices = self._build_prompt_indices(requests_states) -            # Use encode_with_prompts which handles dimension expansion             encoded, encoded_len = self.asr_model.encode_with_prompts(                 processed_signal=processed_signals,                 processed_signal_length=processed_signal_lengths,-                prompt_vectors=prompt_vectors,+                language_prompt_indices=prompt_indices,             )         else:             encoded, encoded_len = self.asr_model.encode(
nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py6 + / 5 −
@@ -52,6 +52,7 @@     get_confidence_utils, ) from nemo.collections.asr.parts.submodules.rnnt_malsd_batched_computer import ModifiedALSDBatchedRNNTComputer+from nemo.collections.asr.parts.utils.prompt_language_utils import resolve_prompt_language from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis from nemo.utils import logging @@ -121,7 +122,10 @@ def init_parameters(self, cfg: DictConfig) -> None:         # Language prompt for multilingual models, taken from the top-level `lang` field.         self.strip_lang_tags = cfg.asr.get("strip_lang_tags", True)         if self.prompt_enabled:-            self.default_language_code = cfg.get("lang", None) or self._resolve_default_language_code()+            # Validated up front, so an unsupported `lang` fails before streaming starts; None selects "auto".+            self.default_language_code = resolve_prompt_language(+                cfg.get("lang", None), self._prompt_config['prompt_dict']+            )             logging.info(f"Multilingual ASR model: conditioning on language '{self.default_language_code}'.")         else:             self.default_language_code = None@@ -250,10 +254,7 @@ def create_state(self, options: ASRRequestOptions) -> CacheAwareRNNTStreamingSta          # Create per-stream prompt index for prompt-enabled models         if self.prompt_enabled:-            lang_code = getattr(new_options, "language_code", None)-            if not isinstance(lang_code, str) or len(lang_code) == 0:-                raise ValueError("Prompt-enabled model requires a valid language_code in request options.")-            prompt_idx = self._resolve_prompt_index(lang_code)+            prompt_idx = self._resolve_prompt_index(getattr(new_options, "language_code", None))             state.set_prompt_index(prompt_idx)          return state
nemo/collections/asr/models/hybrid_rnnt_ctc_bpe_models_prompt.py93 + / 55 −
@@ -16,7 +16,7 @@ import os from dataclasses import dataclass from math import ceil-from typing import Dict, List, Optional, Union+from typing import Any, Dict, List, Optional, Union  import torch from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict@@ -34,6 +34,10 @@ from nemo.collections.asr.parts.preprocessing.segment import ChannelSelectorType from nemo.collections.asr.parts.submodules.ctc_decoding import CTCBPEDecoding, CTCBPEDecodingConfig from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTBPEDecoding+from nemo.collections.asr.parts.utils.prompt_language_utils import (+    AUTO_PROMPT_LANGUAGE,+    resolve_deprecated_prompt_arguments,+) from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config from nemo.core.classes.common import PretrainedModelInfo, typecheck@@ -55,7 +59,13 @@ class HybridRNNTCTCPromptTranscribeConfig(TranscribeConfig):     Configuration for Hybrid RNNT-CTC BPE Model with Prompt Transcription     """ -    target_lang: str = "auto"+    # Language applied to every utterance of the run, e.g. "en-US" or the language-agnostic "auto".+    # None (the default) conditions each utterance on its own manifest language, read from+    # `lang_field`, and uses "auto" for utterances that name none (e.g. a plain list of audio files).+    target_lang: Optional[str] = None++    # Manifest field holding each utterance's language. Only consulted when `target_lang` is None.+    lang_field: str = "lang"   class EncDecHybridRNNTCTCBPEModelWithPrompt(PromptStreamingMixin, EncDecHybridRNNTCTCBPEModel, ASRTranscriptionMixin):@@ -246,6 +256,12 @@ def _setup_dataloader_from_config(self, config: Optional[Dict]):             pin_memory=config.get('pin_memory', False),         ) +    def _transcribe_input_manifest_processing(self, audio_files, temp_dir: str, trcfg) -> Dict[str, Any]:+        """Carry `target_lang` into the dataloader config, so that it conditions every utterance."""+        ds_config = super()._transcribe_input_manifest_processing(audio_files, temp_dir, trcfg)+        ds_config['target_lang'] = getattr(trcfg, 'target_lang', None)+        return ds_config+     def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLoader':         """         Setup function for a temporary data loader which wraps the provided audio file.@@ -258,7 +274,8 @@ def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLo                 Bigger will result in better throughput performance but would use more memory.             temp_dir: (str) A temporary directory where the audio manifest is temporarily                 stored.-            target_lang: (str) target language ID for transcription+            target_lang: (str) language applied to every utterance, or None to use each+                utterance's own language          Returns:             A pytorch DataLoader for the given audio file(s).@@ -270,9 +287,6 @@ def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLo             manifest_filepath = os.path.join(config['temp_dir'], 'manifest.json')             batch_size = min(config['batch_size'], len(config['paths2audio_files'])) -        # Get target language from config-        target_lang = config.get('target_lang', 'en-US')-         dl_config = {             'manifest_filepath': manifest_filepath,             'sample_rate': self.preprocessor._sample_rate,@@ -285,11 +299,15 @@ def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLo             'use_lhotse': config.get('use_lhotse', True),             'use_bucketing': False,             'drop_last': False,+            # "langID" always uses the utterance's language; see the RNNT prompt model.             'initialize_prompt_feature': True,             'prompt_dictionary': self.cfg.model_defaults.get('prompt_dictionary'),             'num_prompts': self.cfg.model_defaults.get('num_prompts', 128),             'subsampling_factor': self.cfg.get('subsampling_factor', 8),-            'default_lang': target_lang,+            'default_prompt_mode': 'langID',+            'target_lang': config.get('target_lang'),+            'default_lang': AUTO_PROMPT_LANGUAGE,+            'lang_field': config.get('lang_field', 'lang'),             'window_stride': self.cfg.preprocessor.get('window_stride', 0.01),         } @@ -315,35 +333,26 @@ def _transcribe_forward(self, batch: tuple[torch.Tensor, ...], trcfg: HybridRNNT             The model's outputs that are processed by `_transcribe_output_processing()`.         """         audio, audio_lens = batch[0], batch[1]-        prompt, prompt_indices = None, None+        language_prompt_features, language_prompt_indices = None, None          if len(batch) >= 5:             prompt_or_indices = batch[4]             if prompt_or_indices.dim() == 1:-                prompt_indices = prompt_or_indices+                language_prompt_indices = prompt_or_indices             else:-                prompt = prompt_or_indices+                language_prompt_features = prompt_or_indices          batch_size = audio.shape[0] -        if prompt is None and prompt_indices is None:-            target_lang = trcfg.target_lang-            prompt_dict = self.cfg.model_defaults.get('prompt_dictionary')--            if not prompt_dict:-                raise ValueError("Prompt dictionary is empty. Cannot create dynamic prompts.")--            if target_lang not in prompt_dict:-                available_keys = list(prompt_dict.keys())-                raise ValueError(-                    f"Unknown target language: '{target_lang}'. Available languages: {available_keys[:10]}{'...' if len(available_keys) > 10 else ''}"-                )--            prompt_id = prompt_dict[target_lang]-            prompt_indices = torch.full((batch_size,), prompt_id, dtype=torch.long, device=audio.device)+        if language_prompt_features is None and language_prompt_indices is None:+            prompt_id = self.resolve_prompt_language(trcfg.target_lang)+            language_prompt_indices = torch.full((batch_size,), prompt_id, dtype=torch.long, device=audio.device)          encoded, encoded_len = self.forward(-            input_signal=audio, input_signal_length=audio_lens, prompt=prompt, prompt_indices=prompt_indices+            input_signal=audio,+            input_signal_length=audio_lens,+            language_prompt_features=language_prompt_features,+            language_prompt_indices=language_prompt_indices,         )          # Prepare output dictionary based on decoder type@@ -395,7 +404,9 @@ def transcribe(             timestamps: (Optional[bool]) timestamps will be returned if set to True as part of hypothesis object             override_config: (Optional[HybridRNNTCTCPromptTranscribeConfig]) override transcription config pre-defined by the user.             **prompt: Optional input to construct the prompts for the model. Accepted formats include:-                target_lang: (str) target language ID for transcription (e.g., "en-US", "de-DE")+                target_lang: (str) language of every utterance (e.g., "en-US", "de", or the language-agnostic+                    "auto"). When omitted, each manifest entry's ``lang_field`` is used, or "auto" if it has none.+                    A language the model does not support raises an error.                 Additional prompt parameters can be passed and will be forwarded to the transcription config.          Returns:@@ -431,8 +442,7 @@ def transcribe(          # Create transcription config if not provided         if override_config is None:-            # Extract target_lang from prompt or use default-            target_lang = prompt.get('target_lang', 'auto')+            defaults = HybridRNNTCTCPromptTranscribeConfig()              trcfg = HybridRNNTCTCPromptTranscribeConfig(                 batch_size=batch_size,@@ -442,7 +452,8 @@ def transcribe(                 augmentor=augmentor,                 verbose=verbose,                 timestamps=timestamps,-                target_lang=target_lang,+                target_lang=prompt.get('target_lang', defaults.target_lang),+                lang_field=prompt.get('lang_field', defaults.lang_field),             )          else:@@ -479,6 +490,9 @@ def input_types(self) -> Optional[Dict[str, NeuralType]]:             "input_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True),             "processed_signal": NeuralType(('B', 'D', 'T'), SpectrogramType(), optional=True),             "processed_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True),+            "language_prompt_features": NeuralType(('B', 'T', 'D'), LabelsType(), optional=True),+            "language_prompt_indices": NeuralType(tuple('B'), LabelsType(), optional=True),+            # Deprecated aliases of language_prompt_features / language_prompt_indices.             "prompt": NeuralType(('B', 'T', 'D'), LabelsType(), optional=True),             "prompt_indices": NeuralType(tuple('B'), LabelsType(), optional=True),         }@@ -497,6 +511,8 @@ def forward(         input_signal_length=None,         processed_signal=None,         processed_signal_length=None,+        language_prompt_features=None,+        language_prompt_indices=None,         prompt=None,         prompt_indices=None,     ):@@ -514,18 +530,24 @@ def forward(                 of shape (B, D, T) that has undergone processing via some DALI preprocessor.             processed_signal_length: Vector of length B, that contains the individual lengths of the                 processed audio sequences.-            prompt: (backward-compatible) Pre-built one-hot prompt tensor of shape [B, T, D].+            language_prompt_features: Pre-built one-hot prompt tensor of shape [B, T, D].                 If provided, used directly (trimmed to encoder length if needed).-            prompt_indices: Tensor of shape [B] containing language ID indices per sample.+            language_prompt_indices: Tensor of shape [B] containing language ID indices per sample.                 The model creates the prompt tensor after encoding using the actual                 encoder output length, guaranteeing no size mismatch.-                Ignored if ``prompt`` is also provided.+                Ignored if ``language_prompt_features`` is also provided.+            prompt: Deprecated alias of ``language_prompt_features``.+            prompt_indices: Deprecated alias of ``language_prompt_indices``.          Returns:             A tuple of 2 elements -             1) The encoded tensor of shape [B, D, T].             2) The lengths of the acoustic sequence after propagation through the encoder, of shape [B].         """+        language_prompt_features, language_prompt_indices = resolve_deprecated_prompt_arguments(+            language_prompt_features, language_prompt_indices, prompt, prompt_indices+        )+         has_input_signal = input_signal is not None and input_signal_length is not None         has_processed_signal = processed_signal is not None and processed_signal_length is not None         if (has_input_signal ^ has_processed_signal) is False:@@ -548,22 +570,29 @@ def forward(         encoded = torch.transpose(encoded, 1, 2)  # B * D * T -> B * T * D          if self.concat:-            if prompt is not None:-                # Backward-compatible path: caller provided a pre-built [B, T, D] one-hot tensor-                if prompt.shape[1] > encoded.shape[1]:-                    prompt = prompt[:, : encoded.shape[1], :]-            elif prompt_indices is not None:-                # New path: build one-hot from per-sample language ID indices+            if language_prompt_features is not None:+                # Caller provided a pre-built [B, T, D] one-hot tensor+                if language_prompt_features.shape[1] > encoded.shape[1]:+                    language_prompt_features = language_prompt_features[:, : encoded.shape[1], :]+            elif language_prompt_indices is not None:+                # Build one-hot from per-sample language ID indices                 batch_size = encoded.shape[0]                 time_steps = encoded.shape[1]                 num_prompts = self.num_prompts-                prompt = torch.zeros(batch_size, time_steps, num_prompts, dtype=encoded.dtype, device=encoded.device)-                prompt.scatter_(2, prompt_indices.view(batch_size, 1, 1).expand(-1, time_steps, -1), 1.0)+                language_prompt_features = torch.zeros(+                    batch_size, time_steps, num_prompts, dtype=encoded.dtype, device=encoded.device+                )+                language_prompt_features.scatter_(+                    2, language_prompt_indices.view(batch_size, 1, 1).expand(-1, time_steps, -1), 1.0+                )             else:-                raise ValueError("Either prompt or prompt_indices must be provided when concat mode is enabled.")+                raise ValueError(+                    "Either language_prompt_features or language_prompt_indices must be provided "+                    "when concat mode is enabled."+                )              out_dtype = encoded.dtype-            concat_enc_states = torch.cat([encoded, prompt], dim=-1)+            concat_enc_states = torch.cat([encoded, language_prompt_features], dim=-1)             encoded = self.prompt_kernel(concat_enc_states).to(out_dtype)          encoded = torch.transpose(encoded, 1, 2)  # B * T * D -> B * D * T@@ -580,17 +609,20 @@ def training_step(self, batch, batch_nb):         signal, signal_len, transcript, transcript_len, prompt_or_indices = batch          # Detect whether batch[4] is old-style prompt [B, T, D] or new-style indices [B]-        prompt, prompt_indices = None, None+        language_prompt_features, language_prompt_indices = None, None         if prompt_or_indices.dim() == 1:-            prompt_indices = prompt_or_indices+            language_prompt_indices = prompt_or_indices         else:-            prompt = prompt_or_indices+            language_prompt_features = prompt_or_indices          if isinstance(batch, DALIOutputs) and batch.has_processed_signal:             encoded, encoded_len = self.forward(processed_signal=signal, processed_signal_length=signal_len)         else:             encoded, encoded_len = self.forward(-                input_signal=signal, input_signal_length=signal_len, prompt=prompt, prompt_indices=prompt_indices+                input_signal=signal,+                input_signal_length=signal_len,+                language_prompt_features=language_prompt_features,+                language_prompt_indices=language_prompt_indices,             )         del signal @@ -699,17 +731,20 @@ def training_step(self, batch, batch_nb):     def predict_step(self, batch, batch_idx, dataloader_idx=0):         signal, signal_len, transcript, transcript_len, prompt_or_indices = batch -        prompt, prompt_indices = None, None+        language_prompt_features, language_prompt_indices = None, None         if prompt_or_indices.dim() == 1:-            prompt_indices = prompt_or_indices+            language_prompt_indices = prompt_or_indices         else:-            prompt = prompt_or_indices+            language_prompt_features = prompt_or_indices          if isinstance(batch, DALIOutputs) and batch.has_processed_signal:             encoded, encoded_len = self.forward(processed_signal=signal, processed_signal_length=signal_len)         else:             encoded, encoded_len = self.forward(-                input_signal=signal, input_signal_length=signal_len, prompt=prompt, prompt_indices=prompt_indices+                input_signal=signal,+                input_signal_length=signal_len,+                language_prompt_features=language_prompt_features,+                language_prompt_indices=language_prompt_indices,             )         del signal @@ -736,17 +771,20 @@ def validation_pass(self, batch, batch_idx, dataloader_idx):          signal, signal_len, transcript, transcript_len, prompt_or_indices = batch -        prompt, prompt_indices = None, None+        language_prompt_features, language_prompt_indices = None, None         if prompt_or_indices.dim() == 1:-            prompt_indices = prompt_or_indices+            language_prompt_indices = prompt_or_indices         else:-            prompt = prompt_or_indices+            language_prompt_features = prompt_or_indices          if isinstance(batch, DALIOutputs) and batch.has_processed_signal:             encoded, encoded_len = self.forward(processed_signal=signal, processed_signal_length=signal_len)         else:             encoded, encoded_len = self.forward(-                input_signal=signal, input_signal_length=signal_len, prompt=prompt, prompt_indices=prompt_indices+                input_signal=signal,+                input_signal_length=signal_len,+                language_prompt_features=language_prompt_features,+                language_prompt_indices=language_prompt_indices,             )         del signal 
nemo/collections/asr/models/rnnt_bpe_models_prompt.py103 + / 47 −
@@ -16,7 +16,7 @@ import os from dataclasses import dataclass from math import ceil-from typing import Dict, List, Optional, Union+from typing import Any, Dict, List, Optional, Union  import torch from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict@@ -32,6 +32,10 @@ from nemo.collections.asr.parts.mixins.transcription import TranscriptionReturnType from nemo.collections.asr.parts.preprocessing.segment import ChannelSelectorType from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTBPEDecoding+from nemo.collections.asr.parts.utils.prompt_language_utils import (+    AUTO_PROMPT_LANGUAGE,+    resolve_deprecated_prompt_arguments,+) from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config from nemo.core.classes.common import PretrainedModelInfo, typecheck@@ -51,7 +55,13 @@ class RNNTPromptTranscribeConfig(TranscribeConfig):     """Transcription configuration for RNNT BPE Model with Prompt conditioning.""" -    target_lang: str = "auto"+    # Language applied to every utterance of the run, e.g. "en", "de-DE" or the language-agnostic+    # "auto". None (the default) conditions each utterance on its own manifest language, read from+    # `lang_field`, and uses "auto" for utterances that name none (e.g. a plain list of audio files).+    target_lang: Optional[str] = None++    # Manifest field holding each utterance's language. Only consulted when `target_lang` is None.+    lang_field: str = "lang"   class EncDecRNNTBPEModelWithPrompt(PromptStreamingMixin, EncDecRNNTBPEModel, ASRTranscriptionMixin):@@ -61,6 +71,12 @@ class EncDecRNNTBPEModelWithPrompt(PromptStreamingMixin, EncDecRNNTBPEModel, ASR     cache-aware streaming model.  The prompt mechanism concatenates a language-ID     one-hot vector to the encoder output and projects back to the original     dimension, allowing the decoder to condition on the target language.++    The class also hosts multilingual ("unified") checkpoints that share this prompt+    mechanism but were trained in dual mode with full-context attention and chunked-limited attention+    with additional right-context. Such checkpoints are transcribed offline or with chunked inference+    (see ``examples/asr/asr_chunked_inference/rnnt``); cache-aware streaming would run+    their encoder outside its trained regime.     """      def __init__(self, cfg: DictConfig, trainer: Trainer = None):@@ -228,6 +244,12 @@ def _setup_dataloader_from_config(self, config: Optional[Dict]):             pin_memory=config.get('pin_memory', False),         ) +    def _transcribe_input_manifest_processing(self, audio_files, temp_dir: str, trcfg) -> Dict[str, Any]:+        """Carry `target_lang` into the dataloader config, so that it conditions every utterance."""+        ds_config = super()._transcribe_input_manifest_processing(audio_files, temp_dir, trcfg)+        ds_config['target_lang'] = getattr(trcfg, 'target_lang', None)+        return ds_config+     def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLoader':         if 'manifest_filepath' in config:             manifest_filepath = config['manifest_filepath']@@ -236,8 +258,6 @@ def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLo             manifest_filepath = os.path.join(config['temp_dir'], 'manifest.json')             batch_size = min(config['batch_size'], len(config['paths2audio_files'])) -        target_lang = config.get('target_lang', 'en-US')-         dl_config = {             'manifest_filepath': manifest_filepath,             'sample_rate': self.preprocessor._sample_rate,@@ -250,11 +270,21 @@ def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLo             'use_lhotse': config.get('use_lhotse', True),             'use_bucketing': False,             'drop_last': False,+            # "langID" always uses the utterance's language. The training default, "unified",+            # substitutes the language-agnostic prompt at random for half the utterances, which+            # would make transcription ignore the requested language. `target_lang`, when set,+            # replaces every utterance's language; `default_lang` covers utterances that name none.             'initialize_prompt_feature': True,             'prompt_dictionary': self.cfg.model_defaults.get('prompt_dictionary'),             'num_prompts': self.cfg.model_defaults.get('num_prompts', 128),+            'default_prompt_mode': 'langID',+            'target_lang': config.get('target_lang'),+            'default_lang': AUTO_PROMPT_LANGUAGE,+            # Which manifest field holds each utterance's language (`gt_lang_attr_name` in+            # transcribe_speech.py). Forwarded so a manifest that names it something other than+            # "lang" still conditions the prompt per utterance.+            'lang_field': config.get('lang_field', 'lang'),             'subsampling_factor': self.cfg.get('subsampling_factor', 8),-            'default_lang': target_lang,             'window_stride': self.cfg.preprocessor.get('window_stride', 0.01),         } @@ -303,7 +333,13 @@ def input_types(self) -> Optional[Dict[str, NeuralType]]:             "input_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True),             "processed_signal": NeuralType(('B', 'D', 'T'), SpectrogramType(), optional=True),             "processed_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True),-            "prompt_indices": NeuralType(tuple('B'), LabelsType()),+            "language_prompt_features": NeuralType(('B', 'T', 'D'), LabelsType(), optional=True),+            "language_prompt_indices": NeuralType(tuple('B'), LabelsType(), optional=True),+            # TODO: `prompt` / `prompt_indices` are deprecated aliases of `language_prompt_features` /+            # `language_prompt_indices`; use the new names. Remove them in next NeMo release, together with the matching+            # forward() parameters and `resolve_deprecated_prompt_arguments`.+            "prompt": NeuralType(('B', 'T', 'D'), LabelsType(), optional=True),+            "prompt_indices": NeuralType(tuple('B'), LabelsType(), optional=True),         }      @property@@ -320,8 +356,34 @@ def forward(         input_signal_length=None,         processed_signal=None,         processed_signal_length=None,+        language_prompt_features=None,+        language_prompt_indices=None,+        prompt=None,         prompt_indices=None,     ):+        """+        Forward pass of the acoustic model, optionally conditioned on a language-ID prompt.++        Args:+            input_signal: Batch of raw audio signals of shape [B, T].+            input_signal_length: Vector of length B with the individual audio lengths.+            processed_signal: Batch of processed audio signals of shape [B, D, T].+            processed_signal_length: Vector of length B with the individual processed lengths.+            language_prompt_features: Pre-built one-hot prompt tensor of shape [B, T, D], trimmed to+                the encoder length when longer. Used by the buffered streaming utilities.+            language_prompt_indices: Tensor of shape [B] with one language-ID index per sample. The+                one-hot is built after encoding from the actual encoder length, so no size mismatch is+                possible. Ignored when ``language_prompt_features`` is also provided.+            prompt: Deprecated alias of ``language_prompt_features``.+            prompt_indices: Deprecated alias of ``language_prompt_indices``.++        Returns:+            A tuple of the encoded tensor of shape [B, D, T] and its lengths of shape [B].+        """+        language_prompt_features, language_prompt_indices = resolve_deprecated_prompt_arguments(+            language_prompt_features, language_prompt_indices, prompt, prompt_indices+        )+         has_input_signal = input_signal is not None and input_signal_length is not None         has_processed_signal = processed_signal is not None and processed_signal_length is not None         if (has_input_signal ^ has_processed_signal) is False:@@ -343,18 +405,28 @@ def forward(         encoded = torch.transpose(encoded, 1, 2)  # B x D x T -> B x T x D          if self.concat:-            if prompt_indices is None:-                raise ValueError("prompt_indices must be provided when concat mode is enabled.")--            batch_size = encoded.shape[0]-            time_steps = encoded.shape[1]-            num_prompts = self.num_prompts--            prompt = torch.zeros(batch_size, time_steps, num_prompts, dtype=encoded.dtype, device=encoded.device)-            prompt.scatter_(2, prompt_indices.view(batch_size, 1, 1).expand(-1, time_steps, -1), 1.0)+            if language_prompt_features is not None:+                if language_prompt_features.shape[1] > encoded.shape[1]:+                    language_prompt_features = language_prompt_features[:, : encoded.shape[1], :]+            elif language_prompt_indices is not None:+                batch_size = encoded.shape[0]+                time_steps = encoded.shape[1]+                num_prompts = self.num_prompts++                language_prompt_features = torch.zeros(+                    batch_size, time_steps, num_prompts, dtype=encoded.dtype, device=encoded.device+                )+                language_prompt_features.scatter_(+                    2, language_prompt_indices.view(batch_size, 1, 1).expand(-1, time_steps, -1), 1.0+                )+            else:+                raise ValueError(+                    "Either language_prompt_features or language_prompt_indices must be provided "+                    "when concat mode is enabled."+                )              out_dtype = encoded.dtype-            concat_enc_states = torch.cat([encoded, prompt], dim=-1)+            concat_enc_states = torch.cat([encoded, language_prompt_features], dim=-1)             encoded = self.prompt_kernel(concat_enc_states).to(out_dtype)          encoded = torch.transpose(encoded, 1, 2)  # B x T x D -> B x D x T@@ -364,13 +436,13 @@ def training_step(self, batch, batch_nb):         if AccessMixin.is_access_enabled(self.model_guid):             AccessMixin.reset_registry(self) -        signal, signal_len, transcript, transcript_len, prompt_indices = batch+        signal, signal_len, transcript, transcript_len, language_prompt_indices = batch          if isinstance(batch, DALIOutputs) and batch.has_processed_signal:             encoded, encoded_len = self.forward(processed_signal=signal, processed_signal_length=signal_len)         else:             encoded, encoded_len = self.forward(-                input_signal=signal, input_signal_length=signal_len, prompt_indices=prompt_indices+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=language_prompt_indices             )         del signal @@ -446,13 +518,13 @@ def training_step(self, batch, batch_nb):         return {'loss': loss_value}      def validation_pass(self, batch, batch_idx, dataloader_idx=0):-        signal, signal_len, transcript, transcript_len, prompt_indices = batch+        signal, signal_len, transcript, transcript_len, language_prompt_indices = batch          if isinstance(batch, DALIOutputs) and batch.has_processed_signal:             encoded, encoded_len = self.forward(processed_signal=signal, processed_signal_length=signal_len)         else:             encoded, encoded_len = self.forward(-                input_signal=signal, input_signal_length=signal_len, prompt_indices=prompt_indices+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=language_prompt_indices             )         del signal @@ -527,13 +599,13 @@ def test_step(self, batch, batch_idx, dataloader_idx=0):         return test_logs      def predict_step(self, batch, batch_idx, dataloader_idx=0):-        signal, signal_len, transcript, transcript_len, prompt_indices = batch+        signal, signal_len, transcript, transcript_len, language_prompt_indices = batch          if isinstance(batch, DALIOutputs) and batch.has_processed_signal:             encoded, encoded_len = self.forward(processed_signal=signal, processed_signal_length=signal_len)         else:             encoded, encoded_len = self.forward(-                input_signal=signal, input_signal_length=signal_len, prompt_indices=prompt_indices+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=language_prompt_indices             )         del signal @@ -548,32 +620,15 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0):      def _transcribe_forward(self, batch, trcfg: RNNTPromptTranscribeConfig) -> dict:         audio, audio_lens = batch[0], batch[1]-        if len(batch) >= 5:-            prompt_indices = batch[4]-        else:-            prompt_indices = None--        batch_size = audio.shape[0]--        if prompt_indices is None:-            target_lang = trcfg.target_lang-            prompt_dict = self.cfg.model_defaults.get('prompt_dictionary')--            if not prompt_dict:-                raise ValueError("Prompt dictionary is empty. Cannot create dynamic prompts.")--            if target_lang not in prompt_dict:-                available_keys = list(prompt_dict.keys())-                raise ValueError(-                    f"Unknown target language: '{target_lang}'. "-                    f"Available languages: {available_keys[:10]}{'...' if len(available_keys) > 10 else ''}"-                )+        language_prompt_indices = batch[4] if len(batch) >= 5 else None -            prompt_id = prompt_dict[target_lang]-            prompt_indices = torch.full((batch_size,), prompt_id, dtype=torch.long, device=audio.device)+        # Batches built from audio tensors carry no per-utterance indices.+        if language_prompt_indices is None:+            prompt_id = self.resolve_prompt_language(trcfg.target_lang)+            language_prompt_indices = torch.full((audio.shape[0],), prompt_id, dtype=torch.long, device=audio.device)          encoded, encoded_len = self.forward(-            input_signal=audio, input_signal_length=audio_lens, prompt_indices=prompt_indices+            input_signal=audio, input_signal_length=audio_lens, language_prompt_indices=language_prompt_indices         )          return dict(encoded=encoded, encoded_len=encoded_len)@@ -607,7 +662,7 @@ def transcribe(             self.change_decoding_strategy(decoding_cfg, verbose=False)          if override_config is None:-            target_lang = prompt.get('target_lang', 'auto')+            defaults = RNNTPromptTranscribeConfig()              trcfg = RNNTPromptTranscribeConfig(                 batch_size=batch_size,@@ -617,7 +672,8 @@ def transcribe(                 augmentor=augmentor,                 verbose=verbose,                 timestamps=timestamps,-                target_lang=target_lang,+                target_lang=prompt.get('target_lang', defaults.target_lang),+                lang_field=prompt.get('lang_field', defaults.lang_field),             )         else:             if not isinstance(override_config, RNNTPromptTranscribeConfig):
nemo/collections/asr/parts/mixins/mixins.py22 + / 11 −
@@ -18,7 +18,7 @@ import tarfile from abc import ABC, abstractmethod from pathlib import PurePosixPath-from typing import List+from typing import List, Optional  import torch from omegaconf import DictConfig, OmegaConf, open_dict@@ -28,6 +28,7 @@ from nemo.collections.asr.parts.mixins.asr_adapter_mixins import ASRAdapterModelMixin from nemo.collections.asr.parts.mixins.streaming import StreamingEncoder from nemo.collections.asr.parts.utils import asr_module_utils+from nemo.collections.asr.parts.utils.prompt_language_utils import resolve_prompt_index from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis from nemo.collections.asr.parts.utils.tokenizer_utils import (     extract_capitalized_tokens_from_vocab,@@ -969,7 +970,24 @@ def initialize_prompt_feature(self):             torch.nn.Linear(proj_out_size * 2, proj_out_size),         ) -    def set_inference_prompt(self, target_lang: str):+    def resolve_prompt_language(self, target_lang: Optional[str]) -> int:+        """Map a language name to its prompt index.++        Accepts an exact ``prompt_dictionary`` key (``"de-DE"``, ``"de"``) or a bare language code+        that matches exactly one locale key. None selects the language-agnostic ``"auto"`` prompt.++        Args:+            target_lang: Language name, or None for the language-agnostic prompt.++        Returns:+            (int) Index into the prompt one-hot vector.++        Raises:+            ValueError: If the language is not supported by the model, or is ambiguous.+        """+        return resolve_prompt_index(target_lang, self.cfg.model_defaults.get('prompt_dictionary') or {})++    def set_inference_prompt(self, target_lang: Optional[str]):         """         Set the language prompt for streaming inference. @@ -978,17 +996,10 @@ def set_inference_prompt(self, target_lang: str):         ``change_decoding_strategy``.          Args:-            target_lang: A key from the model's ``prompt_dictionary``+            target_lang: Language name, resolved by :meth:`resolve_prompt_language`                          (e.g. ``"en-US"``, ``"auto"``).         """-        prompt_dict = self.cfg.model_defaults.get('prompt_dictionary', {})-        if target_lang not in prompt_dict:-            available = list(prompt_dict.keys())-            raise ValueError(-                f"Unknown target language '{target_lang}'. "-                f"Available: {available[:20]}{'...' if len(available) > 20 else ''}"-            )-        self._inference_prompt_index = prompt_dict[target_lang]+        self._inference_prompt_index = self.resolve_prompt_language(target_lang)         logging.info(f"Inference prompt set to '{target_lang}' (index {self._inference_prompt_index})")      def _apply_prompt_to_encoded(self, encoded: Tensor) -> Tensor:
nemo/collections/asr/parts/utils/prompt_language_utils.pyadded130 + / 0 −
@@ -0,0 +1,130 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.+# SPDX-License-Identifier: Apache-2.0+#+# 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.++"""Language-ID prompt resolution shared by prompt-conditioned ASR models, datasets and pipelines."""++from typing import Any, Mapping, Optional, Tuple++from nemo.utils import logging+from nemo.utils.nemo_logging import LogMode++__all__ = [+    'AUTO_PROMPT_LANGUAGE',+    'resolve_prompt_language',+    'resolve_prompt_index',+    'resolve_deprecated_prompt_arguments',+]++# The language-agnostic prompt, used when no language is given.+AUTO_PROMPT_LANGUAGE = "auto"+++def resolve_prompt_language(language: Optional[str], prompt_dictionary: Mapping[str, int]) -> str:+    """Map a language name to its key in a model's prompt dictionary.++    Resolution rules:+        * ``None`` selects the language-agnostic ``"auto"`` prompt.+        * An exact key (``"de-DE"``, ``"de"``) is used as is.+        * A bare language code (``"de"``) matches a locale key (``"de-DE"``) only when exactly one+          key has that code.+        * Anything else raises, so an unsupported language is never silently replaced.++    Args:+        language: Language name, or None for the language-agnostic prompt.+        prompt_dictionary: Mapping from language name to prompt index.++    Returns:+        (str) A key of ``prompt_dictionary``.++    Raises:+        ValueError: If the dictionary is empty, or the language is unsupported or ambiguous.+    """+    if not prompt_dictionary:+        raise ValueError("The model's prompt dictionary is empty, so no language prompt can be selected.")++    if language is None:+        if AUTO_PROMPT_LANGUAGE not in prompt_dictionary:+            raise ValueError(+                f"No language was given and this model has no '{AUTO_PROMPT_LANGUAGE}' prompt. "+                f"Pass one of the supported languages: {_format_languages(prompt_dictionary)}"+            )+        return AUTO_PROMPT_LANGUAGE++    if language in prompt_dictionary:+        return language++    prefix = f"{language.lower()}-"+    matches = sorted(key for key in prompt_dictionary if key.lower().startswith(prefix))+    if len(matches) == 1:+        return matches[0]+    if len(matches) > 1:+        raise ValueError(f"Language '{language}' is ambiguous for this model: it matches {matches}. Pass one of them.")++    message = f"Language '{language}' is not supported by this model. Supported languages: " + _format_languages(+        prompt_dictionary+    )+    if AUTO_PROMPT_LANGUAGE in prompt_dictionary:+        message += (+            f". You can try the language-agnostic '{AUTO_PROMPT_LANGUAGE}' prompt instead, "+            "but its performance on unsupported languages is not guaranteed."+        )+    raise ValueError(message)+++def resolve_prompt_index(language: Optional[str], prompt_dictionary: Mapping[str, int]) -> int:+    """Map a language name to its prompt index; see :func:`resolve_prompt_language` for the rules."""+    return prompt_dictionary[resolve_prompt_language(language, prompt_dictionary)]+++def resolve_deprecated_prompt_arguments(+    language_prompt_features: Any, language_prompt_indices: Any, prompt: Any, prompt_indices: Any+) -> Tuple[Any, Any]:+    """Fold the deprecated ``prompt`` / ``prompt_indices`` forward arguments into their new names.++    Args:+        language_prompt_features: Value passed as ``language_prompt_features``.+        language_prompt_indices: Value passed as ``language_prompt_indices``.+        prompt: Value passed as the deprecated ``prompt``.+        prompt_indices: Value passed as the deprecated ``prompt_indices``.++    Returns:+        (tuple) ``(language_prompt_features, language_prompt_indices)``.++    Raises:+        ValueError: If a deprecated argument and its replacement are both given.+    """+    return (+        _resolve_deprecated_argument(language_prompt_features, 'language_prompt_features', prompt, 'prompt'),+        _resolve_deprecated_argument(+            language_prompt_indices, 'language_prompt_indices', prompt_indices, 'prompt_indices'+        ),+    )+++def _resolve_deprecated_argument(value: Any, name: str, deprecated_value: Any, deprecated_name: str) -> Any:+    if deprecated_value is None:+        return value+    if value is not None:+        raise ValueError(f"Pass either `{name}` or the deprecated `{deprecated_name}`, not both.")+    logging.warning(+        f"The `{deprecated_name}` argument of forward() is deprecated and will be removed in a future release. "+        f"Use `{name}` instead.",+        mode=LogMode.ONCE,+    )+    return deprecated_value+++def _format_languages(prompt_dictionary: Mapping[str, int]) -> str:+    return ", ".join(sorted(key for key in prompt_dictionary if key != AUTO_PROMPT_LANGUAGE))
nemo/collections/asr/parts/utils/streaming_utils.py1 + / 1 −
@@ -1220,7 +1220,7 @@ def _get_batch_preds(self):             encoded, encoded_len = self.asr_model.forward(                 processed_signal=feat_signal,                 processed_signal_length=feat_signal_len,-                prompt=prompt_tensor,+                language_prompt_features=prompt_tensor,             )         else:             encoded, encoded_len = self.asr_model.forward(
tests/collections/asr/test_asr_hybrid_rnnt_ctc_model_bpe_prompt.py36 + / 0 −
@@ -279,6 +279,42 @@ def test_forward_prompt_vs_indices_equivalence(self, hybrid_asr_model_with_promp         assert out_from_prompt.shape == out_from_indices.shape         assert torch.max(torch.abs(out_from_prompt - out_from_indices)) <= 1e-6 +    @pytest.mark.unit+    def test_language_prompt_names_match_deprecated_names(self, hybrid_asr_model_with_prompt):+        """``language_prompt_*`` are the primary names; ``prompt`` / ``prompt_indices`` are deprecated aliases."""+        hybrid_asr_model_with_prompt = hybrid_asr_model_with_prompt.eval()+        hybrid_asr_model_with_prompt.preprocessor.featurizer.dither = 0.0+        hybrid_asr_model_with_prompt.preprocessor.featurizer.pad_to = 0++        input_signal = torch.randn(size=(2, 512))+        length = torch.tensor([512, 400])+        prompt_indices = torch.tensor([0, 3], dtype=torch.long)++        with torch.no_grad():+            from_new, _ = hybrid_asr_model_with_prompt.forward(+                input_signal=input_signal, input_signal_length=length, language_prompt_indices=prompt_indices+            )+            from_deprecated, _ = hybrid_asr_model_with_prompt.forward(+                input_signal=input_signal, input_signal_length=length, prompt_indices=prompt_indices+            )+            time_steps = from_new.shape[2]+            prompt_one_hot = torch.zeros(2, time_steps, hybrid_asr_model_with_prompt.num_prompts)+            prompt_one_hot.scatter_(2, prompt_indices.view(2, 1, 1).expand(-1, time_steps, -1), 1.0)+            from_features, _ = hybrid_asr_model_with_prompt.forward(+                input_signal=input_signal, input_signal_length=length, language_prompt_features=prompt_one_hot+            )++            with pytest.raises(ValueError, match="not both"):+                hybrid_asr_model_with_prompt.forward(+                    input_signal=input_signal,+                    input_signal_length=length,+                    language_prompt_indices=prompt_indices,+                    prompt_indices=prompt_indices,+                )++        assert torch.equal(from_new, from_deprecated)+        assert torch.max(torch.abs(from_new - from_features)) <= 1e-6+     @pytest.mark.unit     def test_predict_step_with_prompt(self, hybrid_asr_model_with_prompt):         """Exercise the legacy ``.dim() == 3`` branch of predict_step (3D one-hot prompt in the batch)."""
tests/collections/asr/test_asr_rnnt_prompt_language.pyadded649 + / 0 −
@@ -0,0 +1,649 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.+# SPDX-License-Identifier: Apache-2.0+#+# 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.++"""Language-ID prompt conditioning in :class:`EncDecRNNTBPEModelWithPrompt`.++Covers the shared language resolver used by the model, the prompt-index dataset and the streaming+pipelines, and the ``prompt`` / ``prompt_indices`` forward arguments.+"""++import os+from functools import partial+from types import SimpleNamespace++import pytest+import torch+from omegaconf import DictConfig++from nemo.collections.asr.data.audio_to_text_lhotse_prompt_index import LhotseSpeechToTextBpeDatasetWithPromptIndex+from nemo.collections.asr.inference.model_wrappers.rnnt_inference_wrapper import RNNTInferenceWrapper+from nemo.collections.asr.inference.pipelines.base_pipeline import BasePipeline+from nemo.collections.asr.models import EncDecRNNTBPEModel, EncDecRNNTModel+from nemo.collections.asr.models.rnnt_bpe_models_prompt import EncDecRNNTBPEModelWithPrompt, RNNTPromptTranscribeConfig+from nemo.collections.asr.parts.utils import prompt_language_utils+from nemo.collections.asr.parts.utils.prompt_language_utils import resolve_prompt_index, resolve_prompt_language+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ as NUMBA_MINIMUM_VERSION+from nemo.core.utils.numba_utils import NUMBA_INSTALLATION_MESSAGE, numba_cpu_is_supported++ENC_HIDDEN = 64+NUM_PROMPTS = 8+PROMPT_DICTIONARY = {'en-US': 0, 'en-GB': 1, 'de-DE': 3, 'ja-JP': 5, 'auto': 6, 'fr': 7}++pytestmark = pytest.mark.skipif(+    not numba_cpu_is_supported(NUMBA_MINIMUM_VERSION),+    reason=f'RNNTLoss has not been compiled with appropriate numba version. {NUMBA_INSTALLATION_MESSAGE}',+)+++def build_config(tokenizer_dir: str, prompt_enabled: bool, prompt_dictionary=None) -> DictConfig:+    """Build a tiny prompt-conditioned RNNT BPE config, optionally with prompting enabled."""+    model_defaults = {'enc_hidden': ENC_HIDDEN, 'pred_hidden': 32}+    if prompt_enabled:+        model_defaults.update(+            {+                'initialize_prompt_feature': True,+                'num_prompts': NUM_PROMPTS,+                'prompt_dictionary': PROMPT_DICTIONARY if prompt_dictionary is None else prompt_dictionary,+            }+        )++    return DictConfig(+        {+            'preprocessor': DictConfig(+                {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': {}}+            ),+            'model_defaults': DictConfig(model_defaults),+            'encoder': DictConfig(+                {+                    'cls': 'nemo.collections.asr.modules.ConvASREncoder',+                    'params': {+                        'feat_in': 64,+                        'activation': 'relu',+                        'conv_mask': True,+                        'jasper': [+                            {+                                'filters': ENC_HIDDEN,+                                'repeat': 1,+                                'kernel': [1],+                                'stride': [1],+                                'dilation': [1],+                                'dropout': 0.0,+                                'residual': False,+                                'separable': True,+                                'se': True,+                                'se_context_size': -1,+                            }+                        ],+                    },+                }+            ),+            'decoder': DictConfig(+                {+                    '_target_': 'nemo.collections.asr.modules.RNNTDecoder',+                    'prednet': {'pred_hidden': 32, 'pred_rnn_layers': 1},+                }+            ),+            'joint': DictConfig(+                {+                    '_target_': 'nemo.collections.asr.modules.RNNTJoint',+                    'jointnet': {'joint_hidden': 32, 'activation': 'relu'},+                }+            ),+            'tokenizer': DictConfig(+                {'dir': os.path.join(tokenizer_dir, 'asr', 'tokenizers', 'an4_wpe_128'), 'type': 'wpe'}+            ),+            'decoding': DictConfig({'strategy': 'greedy_batch', 'greedy': {'max_symbols': 5}}),+            'loss': DictConfig({'loss_name': 'default'}),+        }+    )+++@pytest.fixture()+def build_prompt_model(test_data_dir):+    """Factory for tiny :class:`EncDecRNNTBPEModelWithPrompt` instances."""++    def _build(prompt_enabled=True, prompt_dictionary=None):+        cfg = build_config(test_data_dir, prompt_enabled, prompt_dictionary)+        # eval() disables the preprocessor's dither, so two forward passes over the same audio are+        # comparable; several tests below assert exact agreement between prompt paths.+        return EncDecRNNTBPEModelWithPrompt(cfg=cfg).eval()++    return _build+++@pytest.fixture()+def prompt_model(build_prompt_model):+    """A model with language-ID prompt conditioning enabled."""+    return build_prompt_model(prompt_enabled=True)+++@pytest.fixture()+def unconditioned_model(build_prompt_model):+    """The same class instantiated from a config that does not enable prompting."""+    return build_prompt_model(prompt_enabled=False)+++def make_audio(batch_size=2, samples=4000):+    """A deterministic audio batch and its lengths."""+    generator = torch.Generator().manual_seed(0)+    signal = torch.randn(batch_size, samples, generator=generator)+    return signal, torch.full((batch_size,), samples, dtype=torch.long)+++class TestForwardPromptArguments:+    """``forward`` accepts language indices or a pre-built one-hot, mirroring the hybrid variant."""++    @pytest.mark.unit+    def test_forward_declares_new_and_deprecated_names(self, prompt_model):+        """`typecheck` rejects undeclared keywords, so the deprecated names must stay declared."""+        assert {'language_prompt_features', 'language_prompt_indices', 'prompt', 'prompt_indices'} <= set(+            prompt_model.input_types+        )++    @pytest.mark.unit+    @pytest.mark.parametrize(+        "deprecated_name, new_name",+        [('prompt_indices', 'language_prompt_indices'), ('prompt', 'language_prompt_features')],+    )+    def test_deprecated_name_matches_new_name_and_warns(self, prompt_model, monkeypatch, deprecated_name, new_name):+        warnings = []+        monkeypatch.setattr(+            prompt_language_utils.logging, 'warning', lambda message, *args, **kwargs: warnings.append(message)+        )+        signal, signal_len = make_audio()+        with torch.no_grad():+            reference, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=torch.tensor([3, 3])+            )+            value = torch.tensor([3, 3])+            if deprecated_name == 'prompt':+                value = torch.zeros(2, reference.shape[2], NUM_PROMPTS)+                value[:, :, 3] = 1.0+            new, _ = prompt_model.forward(input_signal=signal, input_signal_length=signal_len, **{new_name: value})+            assert not warnings+            deprecated, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, **{deprecated_name: value}+            )++        assert torch.allclose(new, reference)+        assert torch.allclose(deprecated, reference)+        assert len(warnings) == 1+        assert f"`{deprecated_name}`" in warnings[0] and f"`{new_name}`" in warnings[0]++    @pytest.mark.unit+    @pytest.mark.parametrize(+        "deprecated_name, new_name",+        [('prompt_indices', 'language_prompt_indices'), ('prompt', 'language_prompt_features')],+    )+    def test_deprecated_and_new_name_together_raise(self, prompt_model, deprecated_name, new_name):+        signal, signal_len = make_audio()+        value = torch.zeros(2, 4, NUM_PROMPTS) if deprecated_name == 'prompt' else torch.tensor([3, 3])+        with pytest.raises(ValueError, match=f"either `{new_name}` or the deprecated `{deprecated_name}`"):+            prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, **{new_name: value, deprecated_name: value}+            )++    @pytest.mark.unit+    def test_prompt_indices_condition_the_encoder(self, prompt_model):+        signal, signal_len = make_audio()+        with torch.no_grad():+            de, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=torch.tensor([3, 3])+            )+            ja, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=torch.tensor([5, 5])+            )++        assert not torch.allclose(de, ja), "different languages must produce different encoder output"++    @pytest.mark.unit+    def test_prebuilt_prompt_matches_indices(self, prompt_model):+        """The `prompt` path must agree with the `prompt_indices` path."""+        signal, signal_len = make_audio()+        with torch.no_grad():+            from_indices, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=torch.tensor([3, 3])+            )+            time_steps = from_indices.shape[2]+            one_hot = torch.zeros(2, time_steps, NUM_PROMPTS)+            one_hot[:, :, 3] = 1.0+            from_prompt, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_features=one_hot+            )++        assert torch.allclose(from_indices, from_prompt)++    @pytest.mark.unit+    def test_prompt_longer_than_encoder_output_is_trimmed(self, prompt_model):+        signal, signal_len = make_audio()+        with torch.no_grad():+            reference, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=torch.tensor([3, 3])+            )+            padded = torch.zeros(2, reference.shape[2] + 7, NUM_PROMPTS)+            padded[:, :, 3] = 1.0+            trimmed, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_features=padded+            )++        assert torch.allclose(reference, trimmed)++    @pytest.mark.unit+    def test_per_row_languages_are_independent(self, prompt_model):+        """Each row uses its own language, so a mixed batch matches the single-language runs."""+        signal, signal_len = make_audio()+        with torch.no_grad():+            mixed, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=torch.tensor([3, 5])+            )+            de, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=torch.tensor([3, 3])+            )+            ja, _ = prompt_model.forward(+                input_signal=signal, input_signal_length=signal_len, language_prompt_indices=torch.tensor([5, 5])+            )++        assert torch.allclose(mixed[0], de[0])+        assert torch.allclose(mixed[1], ja[1])++    @pytest.mark.unit+    def test_missing_prompt_raises_when_conditioning_is_enabled(self, prompt_model):+        signal, signal_len = make_audio()+        with pytest.raises(ValueError, match="Either language_prompt_features or language_prompt_indices"):+            prompt_model.forward(input_signal=signal, input_signal_length=signal_len)++    @pytest.mark.unit+    def test_unconditioned_model_needs_no_prompt(self, unconditioned_model):+        """A config that does not enable prompting leaves `forward` unconditioned."""+        assert unconditioned_model.concat is False++        signal, signal_len = make_audio()+        with torch.no_grad():+            encoded, encoded_len = unconditioned_model.forward(input_signal=signal, input_signal_length=signal_len)++        assert encoded.shape[0] == 2+        assert encoded_len.shape == (2,)+++class TestSharedPromptLanguageResolver:+    """The one set of rules used by the models, the prompt-index dataset and the streaming pipelines."""++    @pytest.mark.unit+    def test_exact_dictionary_key(self):+        assert resolve_prompt_language('de-DE', PROMPT_DICTIONARY) == 'de-DE'+        assert resolve_prompt_language('fr', PROMPT_DICTIONARY) == 'fr'++    @pytest.mark.unit+    def test_bare_code_matching_one_locale(self):+        """A manifest carrying "de" must reach the "de-DE" prompt."""+        assert resolve_prompt_index('de', PROMPT_DICTIONARY) == 3+        assert resolve_prompt_index('DE', PROMPT_DICTIONARY) == 3++    @pytest.mark.unit+    def test_none_selects_auto(self):+        assert resolve_prompt_language(None, PROMPT_DICTIONARY) == 'auto'++    @pytest.mark.unit+    def test_none_raises_without_auto(self):+        with pytest.raises(ValueError, match="no 'auto' prompt"):+            resolve_prompt_language(None, {'de-DE': 3})++    @pytest.mark.unit+    def test_ambiguous_code_raises(self):+        """ "en" could mean "en-US" or "en-GB"; picking one silently would be a guess."""+        with pytest.raises(ValueError, match="ambiguous"):+            resolve_prompt_language('en', PROMPT_DICTIONARY)++    @pytest.mark.unit+    def test_unsupported_language_raises_and_suggests_auto(self):+        with pytest.raises(ValueError, match="not supported by this model") as error:+            resolve_prompt_language('kl-KL', PROMPT_DICTIONARY)++        message = str(error.value)+        assert "de-DE" in message, "the supported languages must be listed"+        assert "'auto'" in message and "not guaranteed" in message++    @pytest.mark.unit+    def test_unsupported_language_does_not_suggest_a_missing_auto(self):+        with pytest.raises(ValueError, match="not supported by this model") as error:+            resolve_prompt_language('kl-KL', {'de-DE': 3})++        assert "'auto'" not in str(error.value)++    @pytest.mark.unit+    def test_empty_dictionary_raises(self):+        with pytest.raises(ValueError, match="prompt dictionary is empty"):+            resolve_prompt_language('de-DE', {})+++class TestModelPromptLanguage:+    """The model exposes the shared resolver, with no separate fallback language."""++    @pytest.mark.unit+    def test_resolves_with_the_shared_rules(self, prompt_model):+        assert prompt_model.resolve_prompt_language('de-DE') == 3+        assert prompt_model.resolve_prompt_language('de') == 3+        assert prompt_model.resolve_prompt_language(None) == PROMPT_DICTIONARY['auto']++    @pytest.mark.unit+    def test_unsupported_language_raises(self, prompt_model):+        with pytest.raises(ValueError, match="not supported by this model"):+            prompt_model.resolve_prompt_language('kl-KL')++    @pytest.mark.unit+    def test_no_separate_default_language(self, prompt_model):+        """`auto` is the only language-agnostic prompt, so there is nothing else to configure."""+        assert not hasattr(prompt_model, 'default_prompt_language')++    @pytest.mark.unit+    def test_set_inference_prompt_uses_the_shared_rules(self, prompt_model):+        prompt_model.set_inference_prompt('de')+        assert prompt_model._inference_prompt_index == 3++        with pytest.raises(ValueError, match="not supported by this model"):+            prompt_model.set_inference_prompt('kl-KL')+++class TestTranscribeForward:+    """Offline transcription conditions on the dataloader's indices, else on `target_lang`."""++    @staticmethod+    def _batch(prompt_indices):+        signal, signal_len = make_audio()+        transcript = torch.zeros(2, 4, dtype=torch.long)+        transcript_len = torch.full((2,), 4, dtype=torch.long)+        return signal, signal_len, transcript, transcript_len, prompt_indices++    @pytest.mark.unit+    def test_target_lang_selects_the_language(self, prompt_model):+        batch = self._batch(prompt_indices=None)+        with torch.no_grad():+            de = prompt_model._transcribe_forward(batch, RNNTPromptTranscribeConfig(target_lang='de-DE'))+            ja = prompt_model._transcribe_forward(batch, RNNTPromptTranscribeConfig(target_lang='ja-JP'))++        assert not torch.allclose(de['encoded'], ja['encoded'])++    @pytest.mark.unit+    def test_dataloader_indices_are_used(self, prompt_model):+        """The dataloader already applies `target_lang` or each utterance's own language."""+        with torch.no_grad():+            from_indices = prompt_model._transcribe_forward(+                self._batch(prompt_indices=torch.tensor([3, 3])), RNNTPromptTranscribeConfig(target_lang='ja-JP')+            )+            explicit = prompt_model._transcribe_forward(+                self._batch(prompt_indices=None), RNNTPromptTranscribeConfig(target_lang='de-DE')+            )++        assert torch.allclose(from_indices['encoded'], explicit['encoded'])++    @pytest.mark.unit+    def test_mixed_dataloader_indices_are_honoured_per_row(self, prompt_model):+        """A manifest naming a different language per utterance conditions each row on its own."""+        with torch.no_grad():+            mixed = prompt_model._transcribe_forward(+                self._batch(prompt_indices=torch.tensor([3, 5])), RNNTPromptTranscribeConfig()+            )+            de = prompt_model._transcribe_forward(+                self._batch(prompt_indices=None), RNNTPromptTranscribeConfig(target_lang='de-DE')+            )+            ja = prompt_model._transcribe_forward(+                self._batch(prompt_indices=None), RNNTPromptTranscribeConfig(target_lang='ja-JP')+            )++        assert torch.allclose(mixed['encoded'][0], de['encoded'][0])+        assert torch.allclose(mixed['encoded'][1], ja['encoded'][1])++    @pytest.mark.unit+    def test_auto_is_used_when_nothing_is_given(self, prompt_model):+        with torch.no_grad():+            implicit = prompt_model._transcribe_forward(self._batch(prompt_indices=None), RNNTPromptTranscribeConfig())+            explicit = prompt_model._transcribe_forward(+                self._batch(prompt_indices=None), RNNTPromptTranscribeConfig(target_lang='auto')+            )++        assert torch.allclose(implicit['encoded'], explicit['encoded'])++    @pytest.mark.unit+    def test_unsupported_target_lang_raises(self, prompt_model):+        with pytest.raises(ValueError, match="not supported by this model"):+            prompt_model._transcribe_forward(+                self._batch(prompt_indices=None), RNNTPromptTranscribeConfig(target_lang='kl-KL')+            )++    @pytest.mark.unit+    def test_default_config_reads_the_manifest(self):+        """Offline and streaming share the default: each utterance's own manifest language."""+        assert RNNTPromptTranscribeConfig().target_lang is None+        assert RNNTPromptTranscribeConfig().lang_field == 'lang'+++def make_cut(language, prompt_mode=None):+    """A stand-in cut carrying just the fields prompt resolution reads."""+    custom = None if prompt_mode is None else {'prompt_mode': prompt_mode}+    return SimpleNamespace(id='cut-0', supervisions=[SimpleNamespace(language=language)], custom=custom)+++def make_prompt_dataset(default_prompt_mode='unified', default_lang=None, target_lang=None):+    """The prompt-index dataset, built without audio since only prompt resolution is exercised."""+    cfg = {+        'prompt_dictionary': PROMPT_DICTIONARY,+        'num_prompts': NUM_PROMPTS,+        'default_prompt_mode': default_prompt_mode,+        'default_lang': default_lang,+        'target_lang': target_lang,+    }+    return LhotseSpeechToTextBpeDatasetWithPromptIndex(tokenizer=None, cfg=cfg)+++class TestPromptIndexDatasetLanguage:+    """Per-utterance language resolution in the dataset feeding ``prompt_indices``."""++    @pytest.mark.unit+    def test_langid_mode_uses_the_cut_language(self):+        dataset = make_prompt_dataset(default_prompt_mode='langID')+        assert dataset._get_prompt_index_for_cut(make_cut('de-DE')) == 3+        assert dataset._get_prompt_index_for_cut(make_cut('ja-JP')) == 5++    @pytest.mark.unit+    def test_bare_code_uses_the_shared_rules(self):+        dataset = make_prompt_dataset(default_prompt_mode='langID')+        assert dataset._get_prompt_index_for_cut(make_cut('de')) == 3++    @pytest.mark.unit+    def test_langid_mode_is_deterministic(self):+        """Transcription relies on this: "unified" would substitute `auto` at random instead."""+        dataset = make_prompt_dataset(default_prompt_mode='langID')+        cut = make_cut('de-DE')+        assert {dataset._get_prompt_index_for_cut(cut) for _ in range(50)} == {3}++    @pytest.mark.unit+    def test_missing_language_uses_default_lang(self):+        """A plain list of audio files carries no language, and must still decode."""+        dataset = make_prompt_dataset(default_prompt_mode='langID', default_lang='auto')+        assert dataset._get_prompt_index_for_cut(make_cut(None)) == PROMPT_DICTIONARY['auto']++    @pytest.mark.unit+    def test_missing_language_raises_without_default_lang(self):+        """Training configures no default, so a cut without a language fails loudly."""+        dataset = make_prompt_dataset(default_prompt_mode='langID')+        with pytest.raises(ValueError, match="has no language"):+            dataset._get_prompt_index_for_cut(make_cut(None))++    @pytest.mark.unit+    def test_unsupported_language_raises_even_with_default_lang(self):+        """`default_lang` covers a missing language only; it never replaces an unsupported one."""+        dataset = make_prompt_dataset(default_prompt_mode='langID', default_lang='auto')+        with pytest.raises(ValueError, match="not supported by this model"):+            dataset._get_prompt_index_for_cut(make_cut('kl-KL'))++    @pytest.mark.unit+    def test_target_lang_replaces_every_cut_language(self):+        """A run-wide `target_lang` applies even to cuts whose own language is unsupported."""+        dataset = make_prompt_dataset(default_prompt_mode='langID', target_lang='de-DE')+        assert dataset._get_prompt_index_for_cut(make_cut('ja-JP')) == 3+        assert dataset._get_prompt_index_for_cut(make_cut('kl-KL')) == 3+        assert dataset._get_prompt_index_for_cut(make_cut(None)) == 3++    @pytest.mark.unit+    def test_unsupported_target_lang_raises(self):+        dataset = make_prompt_dataset(default_prompt_mode='langID', target_lang='kl-KL')+        with pytest.raises(ValueError, match="not supported by this model"):+            dataset._get_prompt_index_for_cut(make_cut('de-DE'))++    @pytest.mark.unit+    def test_per_cut_prompt_mode_still_overrides(self):+        """The training tag remains authoritative, so existing data blends are unaffected."""+        dataset = make_prompt_dataset(default_prompt_mode='langID')+        assert dataset._get_prompt_index_for_cut(make_cut('de-DE', prompt_mode='auto')) == dataset.auto_index++    @pytest.mark.unit+    def test_unified_mode_still_randomises(self):+        """The training default is untouched: it mixes the real language with `auto`."""+        dataset = make_prompt_dataset(default_prompt_mode='unified')+        drawn = {dataset._get_prompt_index_for_cut(make_cut('de-DE')) for _ in range(200)}+        assert drawn == {3, dataset.auto_index}+++class TestTranscribeDataloaderLanguage:+    """The transcribe dataloader applies `target_lang`, else each utterance's own language."""++    @staticmethod+    def _capture_dataloader_config(model, monkeypatch, tmp_path, **config):+        captured = {}+        monkeypatch.setattr(model, '_setup_dataloader_from_config', lambda config: captured.update(config))+        manifest = tmp_path / 'manifest.json'+        manifest.write_text('')+        model._setup_transcribe_dataloader({'manifest_filepath': str(manifest), 'batch_size': 2, **config})+        return captured++    @pytest.mark.unit+    def test_dataloader_config(self, prompt_model, monkeypatch, tmp_path):+        captured = self._capture_dataloader_config(prompt_model, monkeypatch, tmp_path)++        assert captured['initialize_prompt_feature'] is True, "per-utterance languages must be read"+        assert captured['default_prompt_mode'] == 'langID', "'unified' would randomise the language"+        assert captured['default_lang'] == 'auto'+        assert captured['target_lang'] is None+        assert captured['prompt_dictionary'] == prompt_model.cfg.model_defaults.prompt_dictionary+        assert captured['lang_field'] == 'lang'++    @pytest.mark.unit+    def test_target_lang_and_lang_field_are_forwarded(self, prompt_model, monkeypatch, tmp_path):+        captured = self._capture_dataloader_config(+            prompt_model, monkeypatch, tmp_path, target_lang='de', lang_field='source_lang'+        )++        assert captured['target_lang'] == 'de'+        assert captured['lang_field'] == 'source_lang'++    @pytest.mark.unit+    def test_manifest_processing_carries_target_lang(self, prompt_model, tmp_path):+        trcfg = RNNTPromptTranscribeConfig(target_lang='de')+        ds_config = prompt_model._transcribe_input_manifest_processing(['a.wav'], str(tmp_path), trcfg)++        assert ds_config['target_lang'] == 'de'+++class TestPipelinePromptLanguage:+    """``BasePipeline`` resolves languages with the same rules as the model."""++    @staticmethod+    def _resolve(language_code, prompt_dict=PROMPT_DICTIONARY):+        stub = SimpleNamespace(_prompt_config={'prompt_dict': prompt_dict})+        return BasePipeline._resolve_prompt_index(stub, language_code)++    @pytest.mark.unit+    def test_exact_and_bare_codes(self):+        assert self._resolve('de-DE') == 3+        assert self._resolve('de') == 3++    @pytest.mark.unit+    @pytest.mark.parametrize("language_code", [None, ''])+    def test_no_language_selects_auto(self, language_code):+        assert self._resolve(language_code) == PROMPT_DICTIONARY['auto']++    @pytest.mark.unit+    def test_unsupported_language_raises(self):+        with pytest.raises(ValueError, match="not supported by this model"):+            self._resolve('kl-KL')++    @pytest.mark.unit+    def test_missing_prompt_config_raises(self):+        with pytest.raises(RuntimeError, match="Prompt configuration is missing"):+            BasePipeline._resolve_prompt_index(SimpleNamespace(_prompt_config=None), 'de-DE')+++class TestEncodeWithPrompts:+    """The buffered wrapper hands the model indices, never a prompt sized by its own length math."""++    @staticmethod+    def _call(processed_signal, captured):+        def model(**kwargs):+            captured.update(kwargs)+            # Subsampling rounds each of its three stages up, so the encoder emits ceil(T / 8).+            time_steps = -(-kwargs['processed_signal'].shape[2] // 8)+            prompt = kwargs.get('language_prompt_features')+            if prompt is not None and prompt.shape[1] < time_steps:+                # The model only trims a prompt that is too long; a short one fails to concatenate.+                raise RuntimeError(f"prompt has {prompt.shape[1]} frames, encoder emitted {time_steps}")+            return torch.zeros(1, 4, time_steps), torch.tensor([time_steps])++        wrapper = SimpleNamespace(+            device=torch.device('cpu'),+            device_str='cpu',+            compute_dtype=torch.float32,+            cast_dtype=torch.float32,+            use_amp=False,+            asr_model=model,+        )+        wrapper.encode = partial(RNNTInferenceWrapper.encode, wrapper)+        return RNNTInferenceWrapper.encode_with_prompts(+            wrapper,+            processed_signal=processed_signal,+            processed_signal_length=torch.tensor([processed_signal.shape[2]]),+            language_prompt_indices=torch.tensor([3]),+        )++    @pytest.mark.unit+    @pytest.mark.parametrize("feature_frames", [768, 700])+    def test_indices_reach_the_model(self, feature_frames):+        """700 frames is the regression: floor(700/8) is 87 while the encoder emits 88."""+        captured = {}+        self._call(torch.zeros(1, 4, feature_frames), captured)++        assert torch.equal(captured['language_prompt_indices'], torch.tensor([3]))+        assert not {'language_prompt_features', 'prompt', 'prompt_indices'} & set(captured)+++class TestBaseClassesAreUntouched:+    """Prompt conditioning stays in the prompt class; stock models gain no prompt surface."""++    @pytest.mark.unit+    @pytest.mark.parametrize(+        "attribute",+        ['concat', 'prompt_kernel', 'resolve_prompt_language', 'set_inference_prompt'],+    )+    def test_plain_rnnt_has_no_prompt_surface(self, attribute):+        assert not hasattr(EncDecRNNTModel, attribute)+        assert not hasattr(EncDecRNNTBPEModel, attribute)++    @pytest.mark.unit+    def test_plain_model_forward_signature_has_no_prompt(self):+        import inspect++        parameters = inspect.signature(EncDecRNNTModel.forward).parameters+        assert not {'language_prompt_features', 'language_prompt_indices', 'prompt', 'prompt_indices'} & set(+            parameters+        )