Add unified prompt model support for multilingual ASR and streaming inference by enas-albasiri · Pull Request #15666 · NVIDIA-NeMo/NeMo

enas-albasiri · 2026-05-06T15:39:37Z

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Add unified prompt architecture for multilingual ASR, enabling language-ID conditioned for both Hybrid RNNT-CTC and RNNT models as well as streaming inference support.

Collection: ASR

Changelog

Add EncDecRNNTBPEModelWithPrompt — RNNT prompt-conditioned model with training script and 600M streaming config
Add LhotseSpeechToTextBpeDatasetWithPromptIndex — index-based Lhotse dataset with per-dataset prompt mode support (langID/auto/unified) via lhotse input_cfg tags
Add backward-compatible forward() in hybrid model accepting both legacy prompt tensors [B, T, D] and new prompt indices [B]
Add set_inference_prompt() + conformer_stream_step() override for streaming inference in both model variants
Add target_lang parameter to speech_to_text_cache_aware_streaming_infer.py (defaults to "auto" for prompt models)
Add strip_lang_tags config in RNNTDecodingConfig to remove <xx-XX> locale tags from decoded output

Usage

# Training with per-dataset prompt modes via lhotse input_cfg:
# In your input_cfg yaml, tag datasets with prompt_mode:
#   - type: nemo_tarred
#     tags: { prompt_mode: langID }    # AST data — always pass lang ID
#   - type: nemo_tarred
#     tags: { prompt_mode: unified }   # ASR data — 50/50 ratio
#   - type: nemo_tarred
#     tags: { prompt_mode: auto }      # Code-switching — always auto

# Streaming inference with target language:
# python speech_to_text_cache_aware_streaming_infer.py \
#     model_path=model.nemo \
#     dataset_manifest=test.json \
#     target_lang=en-US

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

[Y] Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

[Y] New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

- RNNT-only prompt model (EncDecRNNTBPEModelWithPrompt) and training script - RNNT-only streaming config (fastconformer_transducer_bpe_streaming_prompt.yaml) - Index-based dataset (LhotseSpeechToTextBpeDatasetWithPromptIndex) with per-dataset prompt_mode support (langID/auto/unified) via lhotse input_cfg tags - Backward-compatible hybrid model: accepts both old prompt tensors and new prompt_indices with auto-detection - Streaming inference: set_inference_prompt() + conformer_stream_step() override for both hybrid and RNNT-only models, with target_lang support in standard cache-aware streaming inference script - Config-driven strip_lang_tags in RNNT decoding to remove <xx-XX> tags from output - Remove unused docs/TRAINING_GUIDE.md and hybrid streaming config

Signed-off-by: Enas Albasiri <ealbasiri@nvidia.com>

copy-pr-bot · 2026-05-06T15:39:41Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

KunalDhawan

Thanks @enas-albasiri! Could you also share some numbers and validate accuracy of training monolingual/multlingual models with the unified-prompt approach?

Adding @pzelasko for review of LhotseSpeechToTextBpeDatasetWithPromptIndex

KunalDhawan · 2026-05-06T20:12:17Z

+                    )
+                    all_hyp_or_transcribed_texts.append(decoded_out[0])
+            best_hyp = None
+        else:


log_probs is not assigned in the RNNT branch, but is referenced in line 359. Calling conformer_stream_step(..., return_log_probs=True) while cur_decoder == "rnnt" will raise NameError: name 'log_probs' is not defined.

function conformer_stream_step is rewritten based on link, where log_probs only applies for CTC model. RNNT shouldn't bother.

The upstream parallel here is what concerns me, in ASRModuleMixin.conformer_stream_step the if return_log_probs: result.append(log_probs) block sits inside the CTC branch where log_probs is defined. This PR moves the same block outside that branch, after the if isinstance(self, EncDecCTCModel) or ...: ... else: ... split. With the current layout, anything that goes down the RNNT else branch (which is the default, self.cur_decoder == "rnnt") and is called with return_log_probs=True will hit NameError: name 'log_probs' is not defined at line 359, it doesn't need a user to do anything unusual, just pass that flag on an RNNT path.
Three equivalent fixes, any one is fine imo:

Move the if return_log_probs: block back inside the CTC branch (matches upstream exactly).

Initialize log_probs = None at the top of the function and let the RNNT path return None for log_probs.

if return_log_probs and self.cur_decoder == "rnnt": raise NotImplementedError(...) early in the function.

Same shape exists in the RNNT-only variant at rnnt_bpe_models_prompt.py:247, the parameter is in the signature but completely unused; either drop it from the signature or raise NotImplementedError so callers get a clear signal rather than a silently dropped flag.

The entire conformer_stream_step function is removed from both rnnt and hybrid prompt model class. As a replacement, I've added a hook function inside mixin.py which calls model-defined _apply_prompt_to_encoded using the existing conformer_stream_step.

KunalDhawan · 2026-05-06T20:20:30Z


+    # Strip language-ID tags (e.g. <en-US>) from decoded output.
+    # Enable for prompt-conditioned models that emit locale tags after punctuation.
+    strip_lang_tags: bool = True


This is a global default change that touches every existing RNNT model. The dataclass default (True) is also inconsistent with the runtime self.cfg.get('strip_lang_tags', False) (set in line 332 of this file), depending on whether cfg comes from the dataclass vs YAML, user will get different behavior.
Could you set the dataclass default to False and align both, then opt in only via the new YAML config. This is the safer backward-compatible choice.

Set strip_lang_tags to False for inference script.

Thanks! RNNTDecodingConfig.strip_lang_tags: bool = False (line 1872) plus self.cfg.get('strip_lang_tags', False) at line 332 are now consistent, and the streaming-infer script exposes it as an explicit opt-in. Resolving this thread.
One follow-up I'd like to file separately rather than reopening here: the new set_strip_lang_tags method at rnnt_decoding.py:689 only sets self.lang_tag_pattern and not self.strip_lang_tags (the flag the decode path actually reads at line 966), which is why the inference script has to assign both on lines 378–379. I'll leave that as a new inline comment so it doesn't clutter this thread.

You are right. We will come to that point later for sure as it doesn't break the logic for now. There are also other places worth optimization and consolidation, too.

KunalDhawan · 2026-05-06T20:25:13Z

+        trainer=None,
+        validate_access_integrity=True,
+    ):
+        """Delegate to base EncDecRNNTBPEModel to avoid subclass substitution.


This restore_from workaround silently delegates to EncDecRNNTBPEModel.restore_from. As highlighted in your comment, restoring a checkpoint that was saved by EncDecRNNTBPEModelWithPrompt via this method will return a plain EncDecRNNTBPEModel, losing prompt_kernel, set_inference_prompt, and the streaming override. Please fix this.

Losing "prompt_kernel" "set_inference_prompt" "streaming override" will only happen if the model is EncDecRNNTBPEModel class but loaded using EncDecRNNTBPEModelWithPrompt.restore_from(). And it's doesn't matter in this case.

As long as the model is EncDecRNNTBPEModelWithPrompt class, it can be loaded correctly, either loaded using EncDecRNNTBPEModel.restore_from() or EncDecRNNTBPEModelWithPrompt.restore_from(). Or there is other cases where those prompt-specific attr will be lost when loading the prompt model?

KunalDhawan · 2026-05-06T20:38:23Z

+        )
+
+
+class TokenizerWrapper:


This is similar to TokenizerWrapper in NeMo/nemo/collections/common/tokenizers/aggregate_tokenizer.py. Could you please upstream the desired changes to the exisiting TokenizerWrapper and import it here?

Done imported TokenizerWrapper instead.

…script, manual enable by setting stip_lang_tags=True at inference(or training) time. 2. Import TokenizerWrapper in dataloader Signed-off-by: Jinhan <jinhanw@nvidia.com>

pzelasko

Nice work!

pzelasko · 2026-05-15T20:06:17Z

+from nemo.utils import logging
+
+
+class LhotseSpeechToTextBpeDatasetWithPromptIndex(torch.utils.data.Dataset):


Why is a separate class needed for this? Can we replace or add an option to the old class instead?

There is an older version nemo/collections/asr/data/audio_to_text_lhotse_prompt_index.py creating full prompt_tensor from the dataloader, which is heavy-memory. This updated one only returns prompt indices with size [B] (B = batch_size) and model builds the prompt_tensor. This updated one will be used later on. Will work on to remove the older one.

pzelasko · 2026-05-15T20:08:18Z

+        encoded = self.prompt_kernel(torch.cat([encoded, prompt], dim=-1)).to(out_dtype)
+        return encoded.transpose(1, 2)  # (B, T, D) -> (B, D, T)
+
+    def conformer_stream_step(


Should it be encoder_stream_step? Why hardcode the architecture?

Added a hook function in nemo/collections/asr/parts/mixins/mixins.py. Removed conformer_stream_step from both classes.

pzelasko · 2026-05-15T20:09:38Z

        # setting the RNNT decoder as the default one
        self.cur_decoder = "rnnt"

+    # Streaming inference with language-ID prompt


Looks to me like the methods in this section could just as well be a Mixin class PromptableRNNTMixin to avoid duplication in hybrid/non-hybrid model definition

pzelasko · 2026-05-15T20:10:44Z

        # Update the joint fused batch size or disable it entirely if needed.
        self.update_joint_fused_batch_size()
+
+    def set_strip_lang_tags(self, strip_lang_tags: bool):


this method should also accept the regex pattern, with the current one as default if set to None

pzelasko · 2026-05-15T20:13:23Z

+        # Field to use for prompt key (default to 'target_lang')
+        self.prompt_field = cfg.get('prompt_field', 'target_lang')
+
+        self.training_mode = cfg.get('training_mode', True)


Can we remove this and add default_prompt_mode: "inference" instead? One less option.

pzelasko · 2026-05-15T20:14:43Z

+    manifest_filepath: ???
+    sample_rate: ${model.sample_rate}
+    use_lhotse: true
+    shard_manifests: true


this option doesn't exist in lhotse dataloader

pzelasko · 2026-05-15T20:14:52Z

+    min_duration: 0.1
+    is_tarred: true
+    tarred_audio_filepaths: null
+    shuffle_n: 2048


this option doesn't exist in lhotse dataloader

pzelasko · 2026-05-15T20:15:00Z

+    tarred_audio_filepaths: null
+    shuffle_n: 2048
+    slice_length: 100
+    bucketing_strategy: "fully_randomized"


this option doesn't exist in lhotse dataloader

pzelasko · 2026-05-15T20:15:20Z

+    batch_duration: null
+    use_lhotse: true
+    use_bucketing: false
+    max_cuts: 8


this option doesn't exist in lhotse dataloader (use batch_size)

pzelasko · 2026-05-15T20:15:37Z

+    sample_rate: ${model.sample_rate}
+    batch_size: 2
+    shuffle: false
+    use_start_end_token: false


this option doesn't exist in lhotse dataloader (unless you pass it onto the dataset class?)

…nnt prompt model yaml Signed-off-by: Jinhan <jinhanw@nvidia.com>

KunalDhawan · 2026-05-15T21:43:01Z

        # Update the joint fused batch size or disable it entirely if needed.
        self.update_joint_fused_batch_size()
+
+    def set_strip_lang_tags(self, strip_lang_tags: bool):


Currently this method only sets self.lang_tag_pattern (and only when called with True). It does not touch self.strip_lang_tags, which is the flag the decode path actually reads indecode_tokens_to_str_with_strip_punctuation at line 966 (if self.strip_lang_tags: ...). That's why the streaming-infer script has to do this two-step at lines 378–379:

asr_model.decoding.strip_lang_tags = cfg.strip_lang_tags # set the flag asr_model.decoding.set_strip_lang_tags(cfg.strip_lang_tags) # also compile the pattern

Could we either (a) have the method set both the flag and the pattern in one shot:

def set_strip_lang_tags(self, strip_lang_tags: bool): self.strip_lang_tags = strip_lang_tags if strip_lang_tags and not hasattr(self, 'lang_tag_pattern'): self.lang_tag_pattern = re.compile(r'\s*<[a-z]{2}-[A-Z]{2}>')

so the caller becomes asr_model.decoding.set_strip_lang_tags(cfg.strip_lang_tags). Or (b) drop the method entirely, move the regex compile out of the if block in __init__ (line 333) so it's always available, and let callers just assign the attribute? Either way, the regex literal at line 334 and line 692 is currently duplicated, would be good to lift it to a module-level constant.

…ty fallback map - context_manager.update_cache now `.to(dtype=...)`-coerces the encoder's new cache tensors to the pre-allocated cache buffer's dtype before index_copy_. When TRT's io_dtype_overrides doesn't take effect, the engine emits fp32 cache outputs even though apply_fp16_cache_patch allocated fp16 buffers; without the cast we crash on `index_copy_(): self and source expected to have the same dtype`. When dtypes match the cast is a no-op. - Empty _NEMO_TARGET_FALLBACKS: the entry mapping the missing rnnt_bpe_models_prompt.EncDecRNNTBPEModelWithPrompt to the hybrid class is obsolete now that the real class exists on this branch (cherry-picked from PR NVIDIA-NeMo#15666). The peek-target-from-config mechanism stays in place as a safety net for future missing classes.

enas-albasiri added 5 commits May 6, 2026 14:23

Add unified langID support for Hybrid and RNNT

23036c5

Add training guide doc

abbdeb7

fix code style black formatting, flake8 errors, and type hints

6585529

Signed-off-by: Enas Albasiri <ealbasiri@nvidia.com>

fix code style black formatting

2ca55f4

Signed-off-by: Enas Albasiri <ealbasiri@nvidia.com>

github-actions Bot added ASR community-request labels May 6, 2026

KunalDhawan self-requested a review May 6, 2026 20:07

KunalDhawan reviewed May 6, 2026

View reviewed changes

KunalDhawan requested a review from pzelasko May 6, 2026 20:49

Resolve comments: 1. Set strip_lang_tags default False for inference …

299cf1e

…script, manual enable by setting stip_lang_tags=True at inference(or training) time. 2. Import TokenizerWrapper in dataloader Signed-off-by: Jinhan <jinhanw@nvidia.com>

pzelasko requested changes May 15, 2026

View reviewed changes

Resolve critical structural comments and remove un-used config from r…

2d8fcad

…nnt prompt model yaml Signed-off-by: Jinhan <jinhanw@nvidia.com>

KunalDhawan reviewed May 15, 2026

View reviewed changes

		from nemo.utils import logging


		class LhotseSpeechToTextBpeDatasetWithPromptIndex(torch.utils.data.Dataset):

Conversation

enas-albasiri commented May 6, 2026

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

copy-pr-bot Bot commented May 6, 2026

Uh oh!

KunalDhawan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kingformatty May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kingformatty May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kingformatty May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kingformatty May 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pzelasko left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kingformatty May 15, 2026 •

edited

Loading

kingformatty May 13, 2026 •

edited

Loading

kingformatty May 13, 2026 •

edited

Loading

kingformatty May 14, 2026 •

edited

Loading