Adapting transducer greedy decoding (#2975)

younessdkhissi · TParcollet · web-flow · commit 43c8ad1bb8da · 2025-11-02T20:24:25.000Z
Co-authored-by: Parcollet Titouan &lt;parcollet.titouan@gmail.com&gt;
Co-authored-by: Youness Dkhissi &lt;youness.dkhissi@orange.com&gt;
diff --git a/speechbrain/decoders/transducer.py b/speechbrain/decoders/transducer.py
@@ -154,7 +154,11 @@ def forward(self, tn_output):
         return hyps
 
     def transducer_greedy_decode(
-        self, tn_output, hidden_state=None, return_hidden=False
+        self,
+        tn_output,
+        hidden_state=None,
+        return_hidden=False,
+        max_symbols_per_step=5,
     ):
         """Transducer greedy decoder is a greedy decoder over batch which apply Transducer rules:
             1- for each time step in the Transcription Network (TN) output:
@@ -176,6 +180,9 @@ def transducer_greedy_decode(
         return_hidden : bool
             Whether the return tuple should contain an extra 5th element with
             the hidden state at of the last step. See `hidden_state`.
+        max_symbols_per_step : int
+            Maximum number of non-blank symbols to decode per time step. This is
+            useful to avoid infinite loops.
 
         Returns
         -------
@@ -221,42 +228,49 @@ def transducer_greedy_decode(
 
         # For each time step
         for t_step in range(tn_output.size(1)):
-            # do unsqueeze over since tjoint must be have a 4 dim [B,T,U,Hidden]
-            log_probs = self._joint_forward_step(
-                tn_output[:, t_step, :].unsqueeze(1).unsqueeze(1),
-                out_PN.unsqueeze(1),
-            )
-            # Sort outputs at time
-            logp_targets, positions = torch.max(
-                log_probs.squeeze(1).squeeze(1), dim=1
-            )
-            # Batch hidden update
-            have_update_hyp = []
-            for i in range(positions.size(0)):
-                # Update hiddens only if
-                # 1- current prediction is non blank
-                if positions[i].item() != self.blank_id:
-                    hyp["prediction"][i].append(positions[i].item())
-                    hyp["logp_scores"][i] += logp_targets[i]
-                    input_PN[i][0] = positions[i]
-                    have_update_hyp.append(i)
-            if len(have_update_hyp) > 0:
-                # Select sentence to update
-                # And do a forward steps + generated hidden
-                (
-                    selected_input_PN,
-                    selected_hidden,
-                ) = self._get_sentence_to_update(
-                    have_update_hyp, input_PN, hidden
-                )
-                selected_out_PN, selected_hidden = self._forward_PN(
-                    selected_input_PN, self.decode_network_lst, selected_hidden
+            count = 0
+            while count <= max_symbols_per_step:  # avoid infinite loop
+                # do unsqueeze over since tjoint must be have a 4 dim [B,T,U,Hidden]
+                log_probs = self._joint_forward_step(
+                    tn_output[:, t_step, :].unsqueeze(1).unsqueeze(1),
+                    out_PN.unsqueeze(1),
                 )
-                # update hiddens and out_PN
-                out_PN[have_update_hyp] = selected_out_PN
-                hidden = self._update_hiddens(
-                    have_update_hyp, selected_hidden, hidden
+                # Sort outputs at time
+                logp_targets, positions = torch.max(
+                    log_probs.squeeze(1).squeeze(1), dim=1
                 )
+                # Batch hidden update
+                have_update_hyp = []
+                for i in range(positions.size(0)):
+                    # Update hiddens only if
+                    # 1- current prediction is non blank
+                    if positions[i].item() != self.blank_id:
+                        hyp["prediction"][i].append(positions[i].item())
+                        hyp["logp_scores"][i] += logp_targets[i]
+                        input_PN[i][0] = positions[i]
+                        have_update_hyp.append(i)
+                if len(have_update_hyp) > 0:
+                    # Select sentence to update
+                    # And do a forward steps + generated hidden
+                    (
+                        selected_input_PN,
+                        selected_hidden,
+                    ) = self._get_sentence_to_update(
+                        have_update_hyp, input_PN, hidden
+                    )
+                    selected_out_PN, selected_hidden = self._forward_PN(
+                        selected_input_PN,
+                        self.decode_network_lst,
+                        selected_hidden,
+                    )
+                    # update hiddens and out_PN
+                    out_PN[have_update_hyp] = selected_out_PN
+                    hidden = self._update_hiddens(
+                        have_update_hyp, selected_hidden, hidden
+                    )
+                else:
+                    break
+                count += 1
 
         ret = (
             hyp["prediction"],
diff --git a/speechbrain/inference/ASR.py b/speechbrain/inference/ASR.py
@@ -612,6 +612,7 @@ def _get_audio_stream(
             sample_rate=self.audio_normalizer.sample_rate,
             format="fltp",  # torch.float32
             num_channels=1,
+            buffer_chunk_size=-1,  # avoiding the problem of dropping first chunks
         )
 
         for (chunk,) in streamer.stream():

Original file line number	Diff line number	Diff line change
`@@ -612,6 +612,7 @@ def _get_audio_stream(`
`612`	`612`	`sample_rate=self.audio_normalizer.sample_rate,`
`613`	`613`	`format="fltp", # torch.float32`
`614`	`614`	`num_channels=1,`
	`615`	`+ buffer_chunk_size=-1, # avoiding the problem of dropping first chunks`
`615`	`616`	`)`
`616`	`617`
`617`	`618`	`for (chunk,) in streamer.stream():`