NVIDIA-NeMo/Speech · #15940
Add inference support for the Unified multilingual RNNT model with lang prompt
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.EncDecHybridRNNTCTCBPEModelWithPromptdocs/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 statenemo/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+ )