Repository: NVIDIA/NeMo
Pull request #13500: Tdt buffered inference fix
Author: login: hainan-xv
Labels: name: ASR
name: skip-linting
Created: 2025-05-08T16:55:10Z | Merged: 2025-05-12T22:05:48Z

> [!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 ?

Fix issues with TDT buffered inference that results in bad accuracy.

**Collection**: [Note which collection this PR will affect]

ASR

# Changelog

- change alignment representation of TDT to include skipped frames as well
- add a boundary matching algo in the TDT buffered inference procedure.
# Usage

- You can potentially add a usage example below

```
python NeMo/examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py \
               ... \
               merge_algo=tdt
```

# 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**:

- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [x] 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**:

- [ ] New Feature
- [x] 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](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information

- Related to # (issue)



Diff:
diff --git a/examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py b/examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
index c31fa2b9d812..64f7b0ee8260 100644
--- a/examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
+++ b/examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
@@ -23,6 +23,9 @@
 Buffered inference will use large chunk sizes (5-10 seconds) + some additional buffer for context.
 Streaming inference will use small chunk sizes (0.1 to 0.25 seconds) + some additional buffer for context.
 
+Note, currently greedy_batched inferece for TDT is not supported. Decoding strategy will be set to greedy for
+TDT automatically.
+
 # Middle Token merge algorithm
 
 python speech_to_text_buffered_infer_rnnt.py \
@@ -73,6 +76,7 @@
 from nemo.collections.asr.parts.utils.eval_utils import cal_write_wer
 from nemo.collections.asr.parts.utils.streaming_utils import (
     BatchedFrameASRRNNT,
+    BatchedFrameASRTDT,
     LongestCommonSubsequenceBatchedFrameASRRNNT,
 )
 from nemo.collections.asr.parts.utils.transcribe_utils import (
@@ -135,7 +139,10 @@ class TranscriptionConfig:
     stateful_decoding: bool = False  # Whether to perform stateful decoding
 
     # Merge algorithm for transducers
-    merge_algo: Optional[str] = 'middle'  # choices=['middle', 'lcs'], choice of algorithm to apply during inference.
+    # choices=['middle', 'lcs', 'tdt'], choice of algorithm to apply during inference.
+    # if None, we use 'middle' for rnnt and 'tdt' for tdt.
+    merge_algo: Optional[str] = None
+
     lcs_alignment_dir: Optional[str] = None  # Path to a directory to store LCS algo alignments
 
     # Config for word / character error rate calculation
@@ -150,6 +157,8 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:
     """
     Transcribes the input audio and can be used to infer long audio files by chunking
     them into smaller segments.
+    Currently, greedy_batched inferece for TDT is not supported. Decoding strategy
+    will be set to greedy for TDT automatically.
     """
     logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
     torch.set_grad_enabled(False)
@@ -212,9 +221,17 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:
     asr_model.freeze()
     asr_model = asr_model.to(asr_model.device)
 
+    model_is_tdt = hasattr(asr_model.loss, '_loss') and type(asr_model.loss._loss).__name__ == 'TDTLossNumba'
+    if cfg.merge_algo is None:
+        cfg.merge_algo = "tdt" if model_is_tdt else "middle"
+        logging.info(f"merge_algo not specified. We use the default algorithm (middle for rnnt and tdt for tdt).")
+
+    if model_is_tdt and cfg.merge_algo != "tdt":
+        raise ValueError("merge_algo must be 'tdt' for TDT models")
+
     # Change Decoding Config
     with open_dict(cfg.decoding):
-        if cfg.stateful_decoding:
+        if cfg.stateful_decoding or cfg.merge_algo == 'tdt':
             cfg.decoding.strategy = "greedy"
         else:
             cfg.decoding.strategy = "greedy_batch"
@@ -267,6 +284,16 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:
         # Set the LCS algorithm delay.
         frame_asr.lcs_delay = math.floor(((total_buffer - chunk_len)) / model_stride_in_secs)
 
+    elif cfg.merge_algo == 'tdt':
+        frame_asr = BatchedFrameASRTDT(
+            asr_model=asr_model,
+            frame_len=chunk_len,
+            total_buffer=cfg.total_buffer_in_secs,
+            batch_size=cfg.batch_size,
+            max_steps_per_timestep=cfg.max_steps_per_timestep,
+            stateful_decoding=cfg.stateful_decoding,
+        )
+
     else:
         raise ValueError("Invalid choice of merge algorithm for transducer buffered inference.")
 
diff --git a/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py b/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py
index b5495802d2fd..c281fae0be55 100644
--- a/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py
+++ b/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py
@@ -2673,10 +2673,12 @@ def _greedy_decode(
 
             if self.preserve_alignments:
                 # convert Ti-th logits into a torch array
-                hypothesis.alignments.append([])  # blank buffer for next timestep
+                for i in range(skip):
+                    hypothesis.alignments.append([])  # blank buffer until next timestep
 
             if self.preserve_frame_confidence:
-                hypothesis.frame_confidence.append([])  # blank buffer for next timestep
+                for i in range(skip):
+                    hypothesis.frame_confidence.append([])  # blank buffer for next timestep
 
             if symbols_added == self.max_symbols:
                 time_idx += 1
diff --git a/nemo/collections/asr/parts/utils/streaming_utils.py b/nemo/collections/asr/parts/utils/streaming_utils.py
index 46d76dd1c9fa..8982f350a13c 100644
--- a/nemo/collections/asr/parts/utils/streaming_utils.py
+++ b/nemo/collections/asr/parts/utils/streaming_utils.py
@@ -1264,6 +1264,108 @@ def greedy_merge(self, preds):
         return hypothesis
 
 
+class BatchedFrameASRTDT(BatchedFrameASRRNNT):
+    """
+    Batched implementation of FrameBatchASR for TDT models, where the batch dimension is independent audio samples.
+    It's mostly similar to BatchedFrameASRRNNT with special handling of boundary cases due to the frame-skipping
+    resulted by TDT models.
+    """
+
+    def __init__(
+        self,
+        asr_model,
+        frame_len=1.6,
+        total_buffer=4.0,
+        batch_size=32,
+        max_steps_per_timestep: int = 5,
+        stateful_decoding: bool = False,
+        tdt_search_boundary: int = 4,
+    ):
+        '''
+        Args:
+            asr_model: An RNNT model.
+            frame_len: frame's duration, seconds.
+            total_buffer: duration of total audio chunk size, in seconds.
+            batch_size: Number of independent audio samples to process at each step.
+            max_steps_per_timestep: Maximum number of tokens (u) to process per acoustic timestep (t).
+            stateful_decoding: Boolean whether to enable stateful decoding for preservation of state across buffers.
+            tdt_search_boundary: The max number of frames that we search between chunks to match the token at boundary.
+        '''
+        super().__init__(asr_model, frame_len=frame_len, total_buffer=total_buffer, batch_size=batch_size)
+        self.tdt_search_boundary = tdt_search_boundary
+
+    def transcribe(
+        self,
+        tokens_per_chunk: int,
+        delay: int,
+    ):
+        """
+        Performs "middle token" alignment prediction using the buffered audio chunk.
+        """
+        self.infer_logits()
+
+        self.unmerged = [[] for _ in range(self.batch_size)]
+        for idx, alignments in enumerate(self.all_alignments):
+
+            signal_end_idx = self.frame_bufferer.signal_end_index[idx]
+            if signal_end_idx is None:
+                raise ValueError("Signal did not end")
+
+            for a_idx, alignment in enumerate(alignments):
+                if delay == len(alignment):  # chunk size = buffer size
+                    offset = 0
+                else:  # all other cases
+                    offset = 1
+
+                longer_alignment = alignment[
+                    len(alignment)
+                    - offset
+                    - delay
+                    - self.tdt_search_boundary : len(alignment)
+                    - offset
+                    - delay
+                    + tokens_per_chunk
+                ]
+
+                alignment = alignment[
+                    len(alignment) - offset - delay : len(alignment) - offset - delay + tokens_per_chunk
+                ]
+
+                longer_ids, longer_toks = self._alignment_decoder(
+                    longer_alignment, self.asr_model.tokenizer, self.blank_id
+                )
+                ids, _ = self._alignment_decoder(alignment, self.asr_model.tokenizer, self.blank_id)
+
+                if len(longer_ids) > 0 and a_idx < signal_end_idx:
+                    if a_idx == 0 or len(self.unmerged[idx]) == 0:
+                        self.unmerged[idx] = inplace_buffer_merge(
+                            self.unmerged[idx],
+                            ids,
+                            delay,
+                            model=self.asr_model,
+                        )
+                    elif len(self.unmerged[idx]) > 0 and len(longer_toks) > 1:
+                        id_to_match = self.unmerged[idx][-1]
+                        start = min(len(longer_ids) - len(ids), len(longer_ids) - 1)
+                        end = -1
+                        for i in range(start, end, -1):
+                            if longer_ids[i] == id_to_match:
+                                ids = longer_ids[i + 1 :]
+                                break
+
+                        self.unmerged[idx] = inplace_buffer_merge(
+                            self.unmerged[idx],
+                            ids,
+                            delay,
+                            model=self.asr_model,
+                        )
+
+        output = []
+        for idx in range(self.batch_size):
+            output.append(self.greedy_merge(self.unmerged[idx]))
+        return output
+
+
 class LongestCommonSubsequenceBatchedFrameASRRNNT(BatchedFrameASRRNNT):
     """
     Implements a token alignment algorithm for text alignment instead of middle token alignment.


Reviews:
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T17:01:42Z
user: login: hainan-xv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T17:10:07Z
user: login: nithinraok
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T01:45:09Z
user: login: nithinraok
body: Can you add a test case for buffered tdt inference 
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-05-09T01:57:03Z
user: login: hainan-xv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T09:03:11Z
user: login: artbataev
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T10:07:29Z
user: login: hainan-xv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T15:47:46Z
user: login: hainan-xv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T16:23:28Z
user: login: hainan-xv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T16:31:38Z
user: login: artbataev
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T09:10:09Z
user: login: hainan-xv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T13:25:57Z
user: login: nithinraok
body: LGTM! Thanks @hainan-xv 
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-12T15:42:31Z
user: login: artbataev
body: Thanks a lot!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-12T21:48:40Z

Review comments:
diff_hunk: @@ -1263,6 +1263,99 @@
         hypothesis = self.asr_model.tokenizer.ids_to_text(decoded_prediction)
         return hypothesis
 
+class BatchedFrameASRTDT(BatchedFrameASRRNNT):
+    """
+    Batched implementation of FrameBatchASR for TDT models, where the batch dimension is independent audio samples.
+    It's mostly similar to BatchedFrameASRRNNT with special handling of boundary cases due to the frame-skipping
+    resulted by TDT models.
+    """
+
+    def __init__(
+        self,
+        asr_model,
+        frame_len=1.6,
+        total_buffer=4.0,
+        batch_size=32,
+        max_steps_per_timestep: int = 5,
+        stateful_decoding: bool = False,
+        tdt_search_boundary: int = 4,
+    ):
+        '''
+        Args:
+            asr_model: An RNNT model.
+            frame_len: frame's duration, seconds.
+            total_buffer: duration of total audio chunk size, in seconds.
+            batch_size: Number of independent audio samples to process at each step.
+            max_steps_per_timestep: Maximum number of tokens (u) to process per acoustic timestep (t).
+            stateful_decoding: Boolean whether to enable stateful decoding for preservation of state across buffers.
+            tdt_search_boundary: The max number of frames that we search between chunks to match the token at boundary.
+        '''
+        super().__init__(asr_model, frame_len=frame_len, total_buffer=total_buffer, batch_size=batch_size)
+        self.tdt_search_boundary = tdt_search_boundary
+
+    def transcribe(
+        self,
+        tokens_per_chunk: int,
+        delay: int,
+    ):
+        """
+        Performs "middle token" alignment prediction using the buffered audio chunk.
+        """
+        self.infer_logits()
+
+        self.unmerged = [[] for _ in range(self.batch_size)]
+        for idx, alignments in enumerate(self.all_alignments):
+
+            signal_end_idx = self.frame_bufferer.signal_end_index[idx]
+            if signal_end_idx is None:
+                raise ValueError("Signal did not end")
+
+            for a_idx, alignment in enumerate(alignments):
+                if delay == len(alignment):  # chunk size = buffer size
+                    offset = 0
+                else:  # all other cases
+                    offset = 1
+
+                longer_alignment = alignment[
+                    len(alignment) - offset - delay - self.tdt_search_boundary : len(alignment) - offset - delay + tokens_per_chunk
+                ]
+
+                alignment = alignment[
+                    len(alignment) - offset - delay : len(alignment) - offset - delay + tokens_per_chunk
+                ]
+
+                longer_ids, longer_toks = self._alignment_decoder(longer_alignment, self.asr_model.tokenizer, self.blank_id)
+                ids, toks = self._alignment_decoder(alignment, self.asr_model.tokenizer, self.blank_id)
+
+                if len(longer_ids) > 0 and a_idx < signal_end_idx:
+                    if a_idx == 0 or len(self.unmerged[idx]) == 0:
+                        self.unmerged[idx] = inplace_buffer_merge(
+                            self.unmerged[idx],
+                            ids,
+                            delay,
+                            model=self.asr_model,
+                        )
+                    elif len(self.unmerged[idx]) > 0 and len(longer_toks) > 1:
+                        id_to_match = self.unmerged[idx][-1]
+                        start = min(len(longer_ids) - len(ids), len(longer_ids) - 1)
+                        end = -1
+                        for i in range(start, end, -1):
+                            if longer_ids[i] == id_to_match:
+                                ids = longer_ids[i+1:]
+                                toks = longer_toks[i+1:]
path: nemo/collections/asr/parts/utils/streaming_utils.py
user: login: github-advanced-security[bot]
body: ## Variable defined multiple times

This assignment to 'toks' is unnecessary as it is [redefined](1) before this value is used.

[Show more details](https://github.com/NVIDIA/NeMo/security/code-scanning/14849)
created_at: 2025-05-08T17:01:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1263,6 +1263,99 @@
         hypothesis = self.asr_model.tokenizer.ids_to_text(decoded_prediction)
         return hypothesis
 
+class BatchedFrameASRTDT(BatchedFrameASRRNNT):
+    """
+    Batched implementation of FrameBatchASR for TDT models, where the batch dimension is independent audio samples.
+    It's mostly similar to BatchedFrameASRRNNT with special handling of boundary cases due to the frame-skipping
+    resulted by TDT models.
+    """
+
+    def __init__(
+        self,
+        asr_model,
+        frame_len=1.6,
+        total_buffer=4.0,
+        batch_size=32,
+        max_steps_per_timestep: int = 5,
+        stateful_decoding: bool = False,
+        tdt_search_boundary: int = 4,
+    ):
+        '''
+        Args:
+            asr_model: An RNNT model.
+            frame_len: frame's duration, seconds.
+            total_buffer: duration of total audio chunk size, in seconds.
+            batch_size: Number of independent audio samples to process at each step.
+            max_steps_per_timestep: Maximum number of tokens (u) to process per acoustic timestep (t).
+            stateful_decoding: Boolean whether to enable stateful decoding for preservation of state across buffers.
+            tdt_search_boundary: The max number of frames that we search between chunks to match the token at boundary.
+        '''
+        super().__init__(asr_model, frame_len=frame_len, total_buffer=total_buffer, batch_size=batch_size)
+        self.tdt_search_boundary = tdt_search_boundary
+
+    def transcribe(
+        self,
+        tokens_per_chunk: int,
+        delay: int,
+    ):
+        """
+        Performs "middle token" alignment prediction using the buffered audio chunk.
+        """
+        self.infer_logits()
+
+        self.unmerged = [[] for _ in range(self.batch_size)]
+        for idx, alignments in enumerate(self.all_alignments):
+
+            signal_end_idx = self.frame_bufferer.signal_end_index[idx]
+            if signal_end_idx is None:
+                raise ValueError("Signal did not end")
+
+            for a_idx, alignment in enumerate(alignments):
+                if delay == len(alignment):  # chunk size = buffer size
+                    offset = 0
+                else:  # all other cases
+                    offset = 1
+
+                longer_alignment = alignment[
+                    len(alignment) - offset - delay - self.tdt_search_boundary : len(alignment) - offset - delay + tokens_per_chunk
+                ]
+
+                alignment = alignment[
+                    len(alignment) - offset - delay : len(alignment) - offset - delay + tokens_per_chunk
+                ]
+
+                longer_ids, longer_toks = self._alignment_decoder(longer_alignment, self.asr_model.tokenizer, self.blank_id)
+                ids, toks = self._alignment_decoder(alignment, self.asr_model.tokenizer, self.blank_id)
+
+                if len(longer_ids) > 0 and a_idx < signal_end_idx:
+                    if a_idx == 0 or len(self.unmerged[idx]) == 0:
+                        self.unmerged[idx] = inplace_buffer_merge(
+                            self.unmerged[idx],
+                            ids,
+                            delay,
+                            model=self.asr_model,
+                        )
+                    elif len(self.unmerged[idx]) > 0 and len(longer_toks) > 1:
+                        id_to_match = self.unmerged[idx][-1]
+                        start = min(len(longer_ids) - len(ids), len(longer_ids) - 1)
+                        end = -1
+                        for i in range(start, end, -1):
+                            if longer_ids[i] == id_to_match:
+                                ids = longer_ids[i+1:]
+                                toks = longer_toks[i+1:]
path: nemo/collections/asr/parts/utils/streaming_utils.py
user: login: hainan-xv
body: This is a false positive I believe.
created_at: 2025-05-08T17:10:07Z
author_association: COLLABORATOR
diff_hunk: @@ -1263,6 +1263,99 @@
         hypothesis = self.asr_model.tokenizer.ids_to_text(decoded_prediction)
         return hypothesis
 
+class BatchedFrameASRTDT(BatchedFrameASRRNNT):
+    """
+    Batched implementation of FrameBatchASR for TDT models, where the batch dimension is independent audio samples.
+    It's mostly similar to BatchedFrameASRRNNT with special handling of boundary cases due to the frame-skipping
+    resulted by TDT models.
+    """
+
+    def __init__(
+        self,
+        asr_model,
+        frame_len=1.6,
+        total_buffer=4.0,
+        batch_size=32,
+        max_steps_per_timestep: int = 5,
+        stateful_decoding: bool = False,
+        tdt_search_boundary: int = 4,
+    ):
+        '''
+        Args:
+            asr_model: An RNNT model.
+            frame_len: frame's duration, seconds.
+            total_buffer: duration of total audio chunk size, in seconds.
+            batch_size: Number of independent audio samples to process at each step.
+            max_steps_per_timestep: Maximum number of tokens (u) to process per acoustic timestep (t).
+            stateful_decoding: Boolean whether to enable stateful decoding for preservation of state across buffers.
+            tdt_search_boundary: The max number of frames that we search between chunks to match the token at boundary.
+        '''
+        super().__init__(asr_model, frame_len=frame_len, total_buffer=total_buffer, batch_size=batch_size)
+        self.tdt_search_boundary = tdt_search_boundary
+
+    def transcribe(
+        self,
+        tokens_per_chunk: int,
+        delay: int,
+    ):
+        """
+        Performs "middle token" alignment prediction using the buffered audio chunk.
+        """
+        self.infer_logits()
+
+        self.unmerged = [[] for _ in range(self.batch_size)]
+        for idx, alignments in enumerate(self.all_alignments):
+
+            signal_end_idx = self.frame_bufferer.signal_end_index[idx]
+            if signal_end_idx is None:
+                raise ValueError("Signal did not end")
+
+            for a_idx, alignment in enumerate(alignments):
+                if delay == len(alignment):  # chunk size = buffer size
+                    offset = 0
+                else:  # all other cases
+                    offset = 1
+
+                longer_alignment = alignment[
+                    len(alignment) - offset - delay - self.tdt_search_boundary : len(alignment) - offset - delay + tokens_per_chunk
+                ]
+
+                alignment = alignment[
+                    len(alignment) - offset - delay : len(alignment) - offset - delay + tokens_per_chunk
+                ]
+
+                longer_ids, longer_toks = self._alignment_decoder(longer_alignment, self.asr_model.tokenizer, self.blank_id)
+                ids, toks = self._alignment_decoder(alignment, self.asr_model.tokenizer, self.blank_id)
+
+                if len(longer_ids) > 0 and a_idx < signal_end_idx:
+                    if a_idx == 0 or len(self.unmerged[idx]) == 0:
+                        self.unmerged[idx] = inplace_buffer_merge(
+                            self.unmerged[idx],
+                            ids,
+                            delay,
+                            model=self.asr_model,
+                        )
+                    elif len(self.unmerged[idx]) > 0 and len(longer_toks) > 1:
+                        id_to_match = self.unmerged[idx][-1]
+                        start = min(len(longer_ids) - len(ids), len(longer_ids) - 1)
+                        end = -1
+                        for i in range(start, end, -1):
+                            if longer_ids[i] == id_to_match:
+                                ids = longer_ids[i+1:]
+                                toks = longer_toks[i+1:]
path: nemo/collections/asr/parts/utils/streaming_utils.py
user: login: nithinraok
body: I don;t think so. `toks` variable is not used, only gets assigned! 
created_at: 2025-05-09T01:45:09Z
author_association: MEMBER
diff_hunk: @@ -214,7 +217,7 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:
 
     # Change Decoding Config
     with open_dict(cfg.decoding):
-        if cfg.stateful_decoding:
+        if cfg.stateful_decoding or cfg.merge_algo == 'tdt':
path: examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
user: login: nithinraok
body: Add this info to the doc string above and also add to script usage at the top of this script. 
created_at: 2025-05-09T01:52:15Z
author_association: MEMBER
diff_hunk: @@ -1263,6 +1263,99 @@
         hypothesis = self.asr_model.tokenizer.ids_to_text(decoded_prediction)
         return hypothesis
 
+class BatchedFrameASRTDT(BatchedFrameASRRNNT):
+    """
+    Batched implementation of FrameBatchASR for TDT models, where the batch dimension is independent audio samples.
+    It's mostly similar to BatchedFrameASRRNNT with special handling of boundary cases due to the frame-skipping
+    resulted by TDT models.
+    """
+
+    def __init__(
+        self,
+        asr_model,
+        frame_len=1.6,
+        total_buffer=4.0,
+        batch_size=32,
+        max_steps_per_timestep: int = 5,
+        stateful_decoding: bool = False,
+        tdt_search_boundary: int = 4,
+    ):
+        '''
+        Args:
+            asr_model: An RNNT model.
+            frame_len: frame's duration, seconds.
+            total_buffer: duration of total audio chunk size, in seconds.
+            batch_size: Number of independent audio samples to process at each step.
+            max_steps_per_timestep: Maximum number of tokens (u) to process per acoustic timestep (t).
+            stateful_decoding: Boolean whether to enable stateful decoding for preservation of state across buffers.
+            tdt_search_boundary: The max number of frames that we search between chunks to match the token at boundary.
+        '''
+        super().__init__(asr_model, frame_len=frame_len, total_buffer=total_buffer, batch_size=batch_size)
+        self.tdt_search_boundary = tdt_search_boundary
+
+    def transcribe(
+        self,
+        tokens_per_chunk: int,
+        delay: int,
+    ):
+        """
+        Performs "middle token" alignment prediction using the buffered audio chunk.
+        """
+        self.infer_logits()
+
+        self.unmerged = [[] for _ in range(self.batch_size)]
+        for idx, alignments in enumerate(self.all_alignments):
+
+            signal_end_idx = self.frame_bufferer.signal_end_index[idx]
+            if signal_end_idx is None:
+                raise ValueError("Signal did not end")
+
+            for a_idx, alignment in enumerate(alignments):
+                if delay == len(alignment):  # chunk size = buffer size
+                    offset = 0
+                else:  # all other cases
+                    offset = 1
+
+                longer_alignment = alignment[
+                    len(alignment) - offset - delay - self.tdt_search_boundary : len(alignment) - offset - delay + tokens_per_chunk
+                ]
+
+                alignment = alignment[
+                    len(alignment) - offset - delay : len(alignment) - offset - delay + tokens_per_chunk
+                ]
+
+                longer_ids, longer_toks = self._alignment_decoder(longer_alignment, self.asr_model.tokenizer, self.blank_id)
+                ids, toks = self._alignment_decoder(alignment, self.asr_model.tokenizer, self.blank_id)
+
+                if len(longer_ids) > 0 and a_idx < signal_end_idx:
+                    if a_idx == 0 or len(self.unmerged[idx]) == 0:
+                        self.unmerged[idx] = inplace_buffer_merge(
+                            self.unmerged[idx],
+                            ids,
+                            delay,
+                            model=self.asr_model,
+                        )
+                    elif len(self.unmerged[idx]) > 0 and len(longer_toks) > 1:
+                        id_to_match = self.unmerged[idx][-1]
+                        start = min(len(longer_ids) - len(ids), len(longer_ids) - 1)
+                        end = -1
+                        for i in range(start, end, -1):
+                            if longer_ids[i] == id_to_match:
+                                ids = longer_ids[i+1:]
+                                toks = longer_toks[i+1:]
path: nemo/collections/asr/parts/utils/streaming_utils.py
user: login: hainan-xv
body: Ah right you're right.
created_at: 2025-05-09T09:03:10Z
author_association: COLLABORATOR
diff_hunk: @@ -135,7 +136,9 @@ class TranscriptionConfig:
     stateful_decoding: bool = False  # Whether to perform stateful decoding
 
     # Merge algorithm for transducers
-    merge_algo: Optional[str] = 'middle'  # choices=['middle', 'lcs'], choice of algorithm to apply during inference.
+    merge_algo: Optional[str] = (
path: examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
user: login: artbataev
body: Suggestion: set `merge_algo` to `None` (meaning "auto")
In the code use something like: 
```python
if merge_algo is None:
    merge_algo = "tdt" if model.is_tdt() else "middle"
```
created_at: 2025-05-09T10:07:29Z
author_association: COLLABORATOR
diff_hunk: @@ -135,7 +136,9 @@ class TranscriptionConfig:
     stateful_decoding: bool = False  # Whether to perform stateful decoding
 
     # Merge algorithm for transducers
-    merge_algo: Optional[str] = 'middle'  # choices=['middle', 'lcs'], choice of algorithm to apply during inference.
+    merge_algo: Optional[str] = (
path: examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
user: login: hainan-xv
body: Good idea.
created_at: 2025-05-09T15:47:46Z
author_association: COLLABORATOR
diff_hunk: @@ -214,7 +217,7 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:
 
     # Change Decoding Config
     with open_dict(cfg.decoding):
-        if cfg.stateful_decoding:
+        if cfg.stateful_decoding or cfg.merge_algo == 'tdt':
path: examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
user: login: hainan-xv
body: Done.
created_at: 2025-05-09T16:23:28Z
author_association: COLLABORATOR
diff_hunk: @@ -135,7 +136,9 @@ class TranscriptionConfig:
     stateful_decoding: bool = False  # Whether to perform stateful decoding
 
     # Merge algorithm for transducers
-    merge_algo: Optional[str] = 'middle'  # choices=['middle', 'lcs'], choice of algorithm to apply during inference.
+    merge_algo: Optional[str] = (
path: examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
user: login: hainan-xv
body: Done.
created_at: 2025-05-09T16:31:38Z
author_association: COLLABORATOR
diff_hunk: @@ -135,7 +136,9 @@ class TranscriptionConfig:
     stateful_decoding: bool = False  # Whether to perform stateful decoding
 
     # Merge algorithm for transducers
-    merge_algo: Optional[str] = 'middle'  # choices=['middle', 'lcs'], choice of algorithm to apply during inference.
+    merge_algo: Optional[str] = (
path: examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
user: login: artbataev
body: Thanks! Could you also set default to `None` in the config (here)?
created_at: 2025-05-12T09:10:09Z
author_association: COLLABORATOR
diff_hunk: @@ -135,7 +136,9 @@ class TranscriptionConfig:
     stateful_decoding: bool = False  # Whether to perform stateful decoding
 
     # Merge algorithm for transducers
-    merge_algo: Optional[str] = 'middle'  # choices=['middle', 'lcs'], choice of algorithm to apply during inference.
+    merge_algo: Optional[str] = (
path: examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py
user: login: hainan-xv
body: Done.
created_at: 2025-05-12T13:25:56Z
author_association: COLLABORATOR

PR comments:
user: login: hainan-xv
created_at: 2025-05-08T16:59:16Z
body: @nithinraok can you take a look by any chance? So far I'm not sure it's 100% correct yet, since it's noticeably worse than offline inference in my testing. I'm not sure how much we're expecting the gap to be though... Could you let me know if it works OK on your end?
author_association: COLLABORATOR
user: login: nithinraok
created_at: 2025-05-08T17:01:19Z
body: @weiqingw4ng tested it for few samples and it worked great! I will add some numbers here with `parakeet-tdt-0.6b-v2` on LS dataset in sometime. 
author_association: MEMBER
user: login: nithinraok
created_at: 2025-05-09T01:40:01Z
body: Below table shows evaluation of same size, same data trained parakeet models with various decoders.  Evaluated with greedy (total buffer size, chunk len in sec)

| Decoder | MODE                   | LS-test clean | LS-test other |
|---------|------------------------|---------------|---------------|
|   tdt  | offline                | 1.69          | 3.19          |
| tdt     | greedy (30,10)         | 2.07          | 4.01          |
|   tdt   | greedy (8,0.16)        | 3.23          | 5.19          |
| rnnt    | offline                | 1.82          | 3.47          |
| rnnt    | greedy (30,10)         | 3.09          | 5.14          |
| ctc     | offline                | 2.2           | 3.8           |
| ctc     | greedy (30,10)         | 2.73          | 4.95          |
| ctc     | greedy (8,0.16)  | 5.39          | 7.03          |
author_association: MEMBER
user: login: hainan-xv
created_at: 2025-05-09T16:27:16Z
body: > Can you add a test case for buffered tdt inference

I can't seem to find test code for the existing buffered RNNT inference to base the new test on. Not sure how to do it efficiently... What I can think of is to include a small trained TDT model in the repo and include a short audio recording also to run this, but this will take quite some work... Shall we do this later?
author_association: COLLABORATOR
user: login: nithinraok
created_at: 2025-05-09T18:28:48Z
body: I think something messed up. Now its showing 21 file changes
author_association: MEMBER
user: login: hainan-xv
created_at: 2025-05-10T08:11:22Z
body: > I think something messed up. Now its showing 21 file changes

Not sure what happened back then. I did a pull of the latest main and re-push'ed. It seems it's fixed now.
author_association: COLLABORATOR