From cedfd4624f0aef5531251ebe9f69e8305dff3864 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 5 Jul 2023 11:03:57 +0200
Subject: [PATCH 01/83] Introduce DCT+DCConv logic

---
 .../hparams/conformer_transducer.yaml         |  17 +++
 recipes/LibriSpeech/ASR/transducer/train.py   |  27 ++++
 .../lobes/models/transformer/Conformer.py     | 138 +++++++++++++++++-
 .../models/transformer/TransformerASR.py      |  44 +++++-
 speechbrain/nnet/attention.py                 |  17 +++
 5 files changed, 229 insertions(+), 14 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index dad6b00e26..917cd95c39 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -71,6 +71,23 @@ n_fft: 512
 n_mels: 80
 win_length: 32
 
+# Streaming & dynamic chunk training options
+streaming: False  # controls all DCT & chunk size & left context mechanisms
+
+test_chunk_size: 8
+test_left_context_size: 64
+
+valid_chunk_size: -1
+valid_left_context_size: -1
+
+dynamic_chunk_thresh: 0.6
+dynamic_chunk_min: 8
+dynamic_chunk_max: 32
+
+dynamic_left_context_thresh: 0.75 # TODO: rename all to prob, makes more sense
+dynamic_left_context_min: 16
+dynamic_left_context_max: 64
+
 # Dataloader options
 train_dataloader_opts:
    batch_size: !ref <batch_size>
diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 319edaa7e8..9ad8bb699c 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -69,6 +69,33 @@ def compute_forward(self, batch, stage):
             )
 
         current_epoch = self.hparams.epoch_counter.current
+
+        transformer_chunk_size = -1
+        left_context_chunks = -1
+        if self.hparams.streaming:
+            if stage == sb.Stage.TRAIN:
+                if torch.rand((1, )).item() < self.hparams.dynamic_chunk_thresh:
+                    transformer_chunk_size = torch.randint(
+                        self.hparams.dynamic_chunk_min,
+                        self.hparams.dynamic_chunk_max + 1,
+                        (1, )
+                    ).item()
+                    # print("tfx chunk size", transformer_chunk_size)
+
+                if torch.rand((1, )).item() < self.hparams.dynamic_left_context_thresh:
+                    left_context_chunks = torch.randint(
+                        self.hparams.dynamic_left_context_min,
+                        self.hparams.dynamic_left_context_max + 1,
+                        (1, )
+                    ).item()
+            elif stage == sb.Stage.TEST:
+                transformer_chunk_size = self.hparams.test_chunk_size
+                left_context_chunks = self.hparams.test_left_context_size
+            elif stage == sb.Stage.VALID:
+                transformer_chunk_size = self.hparams.valid_chunk_size
+                left_context_chunks = self.hparams.valid_left_context_size
+
+        # logger.info(f"Batch uses tfx chunk size = {transformer_chunk_size}, frame chunk_size = {chunk_size}")
         feats = self.modules.normalize(feats, wav_lens, epoch=current_epoch)
 
         src = self.modules.CNN(feats)
diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 009b87d932..7a7db13405 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -7,8 +7,10 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from typing import Optional
 import speechbrain as sb
+import math
 import warnings
 
 
@@ -91,28 +93,136 @@ def __init__(
             bias=bias,
         )
 
+        # self.batch_norm = nn.BatchNorm1d(input_size)
+
         self.after_conv = nn.Sequential(
-            nn.LayerNorm(input_size),
+            nn.LayerNorm(input_size),  # should be a BN to match Conformer paper
             activation(),
             # pointwise
             nn.Linear(input_size, input_size, bias=bias),
             nn.Dropout(dropout),
         )
 
-    def forward(self, x, mask=None):
-        """ Processes the input tensor x and returns the output an output tensor"""
+    def _do_conv(self, x, inhibit_padding: bool):
         out = self.layer_norm(x)
         out = out.transpose(1, 2)
         out = self.bottleneck(out)
-        out = self.conv(out)
+
+        if not inhibit_padding:
+            out = self.conv(out)
+        else:
+            # let's keep backwards compat by pointing at the weights from the
+            # already declared Conv1d.
+
+            # we do not need to edit bottleneck as it is pointwise (i.e. time
+            # step by time step), thus, it doesn't need padding along the
+            # time dimension
+            out = F.conv1d(
+                out,
+                weight=self.conv.weight,
+                bias=self.conv.bias,
+                stride=self.conv.stride,
+                padding=0,
+                dilation=self.conv.dilation,
+                groups=out.shape[-2],
+            )
 
         if self.causal:
             # chomp
             out = out[..., : -self.padding]
+
+        # out = self.batch_norm(out)
+
         out = out.transpose(1, 2)
         out = self.after_conv(out)
+        return out
+
+    def forward(self, x, mask=None, chunk_size=-1):
+        """ Processes the input tensor x and returns the output an output tensor"""
+
+        # ref: Dynamic chunk convolution for unified streaming and non-streaming
+        # conformer ASR
+        # https://www.amazon.science/publications/dynamic-chunk-convolution-for-unified-streaming-and-non-streaming-conformer-asr
+        # split the input into chunks of size `chunk_size`, but for each chunk
+        # provide a left context for left chunk dependencies to be possible.
+        if chunk_size >= 1:
+            # chances are chunking+causal is unintended; i don't know where it
+            # may make sense, but if it does to you, feel free to implement it.
+            assert (
+                not self.causal
+            ), "Chunked convolution not supported with causal padding"
+
+            batch_size = x.shape[0]
+            chunk_left_context = self.padding
+
+            chunk_count = int(math.ceil(x.shape[1] / chunk_size))
+
+            if x.shape[1] % chunk_size != 0:
+                final_right_padding = chunk_size - (x.shape[1] % chunk_size)
+            else:
+                final_right_padding = 0
+
+            # compute the left context that can and should be added, for each
+            # chunk. for the first few chunks, we will need to add extra padding
+            applied_left_context = [
+                min(
+                    chunk_left_context,
+                    i * chunk_size,
+                )
+                for i in range(chunk_count)
+            ]
+
+            # build views of chunks with left context (but no 0-padding yet)
+            # the left context effectively becomes "left padding", we do not
+            # want to keep any convolution results centered on the left context
+            out = [
+                x[:,i * chunk_size - applied_left_context[i]:(i + 1) * chunk_size,...]
+                for i in range(chunk_count)
+            ]
+
+            # TODO: experiment around reflect padding, which is difficult
+            # because small chunks have too little time steps to reflect from
+            out = [
+                F.pad(out[i], (
+                    # channel dims, we do not to pad these
+                    0,
+                    0,
+                    # add missing left 0-padding if we lacked left context
+                    chunk_left_context - applied_left_context[i],
+                    # add missing right 0-padding as we disable default padding
+                    # also add missing frames of the rightmost chunk
+                    self.padding + (final_right_padding if i == len(out) - 1 else 0)
+                ))
+                for i in range(len(out))
+            ]
+
+            # we pack together chunks in a single tensor so that we can feed it
+            # to the convolution directly.
+
+            # -> [batch_size, num_chunks, chunk_size + lc + rpad, in_channels]
+            out = torch.stack(out, dim=1)
+
+            # -> [batch_size * num_chunks, chunk_size + lc + rpad, in_channels]
+            out = torch.flatten(out, end_dim=1)
+
+            # -> [batch_size * num_chunks, chunk_size, out_channels]
+            out = self._do_conv(out, inhibit_padding=True)
+
+            # -> [batch_size, num_chunks, chunk_size, out_channels]
+            out = torch.unflatten(out, dim=0, sizes=(batch_size, -1))
+
+            # -> [batch_size, time_steps + extra right padding, out_channels]
+            out = torch.flatten(out, start_dim=1, end_dim=2)
+
+            # -> [batch_size, time_steps, out_channels]
+            if final_right_padding > 0:
+                out = out[:,:-final_right_padding,:]
+        else:
+            out = self._do_conv(x, inhibit_padding=False)
+
         if mask is not None:
             out.masked_fill_(mask, 0.0)
+
         return out
 
 
@@ -231,7 +341,8 @@ def forward(
         x,
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
-        pos_embs: Optional[torch.Tensor] = None,
+        pos_embs: torch.Tensor = None,
+        chunk_size: Optional[int] = None,
     ):
         """
         Arguments
@@ -244,8 +355,14 @@ def forward(
             The mask for the src keys per batch.
         pos_embs: torch.Tensor, torch.nn.Module, optional
             Module or tensor containing the input sequence positional embeddings
+        chunk_size: int, optional
+            Whether to preform convolution chunking to hide future context,
+            useful for chunked conformers in a dynamic chunk training setting
         """
-        conv_mask = None
+        # TODO: cite paper for chunk size
+        # TODO: document left frames
+
+        conv_mask: Optional[torch.Tensor] = None
         if src_key_padding_mask is not None:
             conv_mask = src_key_padding_mask.unsqueeze(-1)
         # ffn module
@@ -253,6 +370,7 @@ def forward(
         # muti-head attention module
         skip = x
         x = self.norm1(x)
+
         x, self_attn = self.mha_layer(
             x,
             x,
@@ -263,7 +381,7 @@ def forward(
         )
         x = x + skip
         # convolution module
-        x = x + self.convolution_module(x, conv_mask)
+        x = x + self.convolution_module(x, conv_mask, chunk_size=chunk_size)
         # ffn module
         x = self.norm2(x + 0.5 * self.ffn_module2(x))
         return x, self_attn
@@ -355,6 +473,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
+        chunk_size: Optional[int] = None,
     ):
         """
         Arguments
@@ -369,8 +488,10 @@ def forward(
             Module or tensor containing the input sequence positional embeddings
             If custom pos_embs are given it needs to have the shape (1, 2*S-1, E)
             where S is the sequence length, and E is the embedding dimension.
+        chunk_size: int, optional
+            Whether to preform convolution chunking to hide future context,
+            useful for chunked conformers in a dynamic chunk training setting
         """
-
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:
                 raise ValueError(
@@ -385,6 +506,7 @@ def forward(
                 src_mask=src_mask,
                 src_key_padding_mask=src_key_padding_mask,
                 pos_embs=pos_embs,
+                chunk_size=chunk_size,
             )
             attention_lst.append(attention)
         output = self.norm(output)
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 1ad0b5c433..7e248acbf9 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -172,6 +172,7 @@ def forward(self, src, tgt, wav_len=None, pad_idx=0):
         pad_idx : int, optional
             The index for <pad> token (default=0).
         """
+        # FIXME: should have chunk_size & left_chunk_context?
 
         # reshpae the src vector to [Batch, Time, Fea] is a 4d vector is given
         if src.ndim == 4:
@@ -229,7 +230,7 @@ def forward(self, src, tgt, wav_len=None, pad_idx=0):
 
         return encoder_out, decoder_out
 
-    def make_masks(self, src, tgt=None, wav_len=None, pad_idx=0):
+    def make_masks(self, src, tgt=None, wav_len=None, pad_idx=0, chunk_size: int = -1, left_context_chunks: int = -1):
         """This method generates the masks for training the transformer model.
 
         Arguments
@@ -246,7 +247,36 @@ def make_masks(self, src, tgt=None, wav_len=None, pad_idx=0):
             abs_len = torch.round(wav_len * src.shape[1])
             src_key_padding_mask = ~length_to_mask(abs_len).bool()
 
-        src_mask = None
+        if chunk_size >= 0 or left_context_chunks >= 0:
+            # wav_len unspecified? make a mask that masks nothing by default
+            # 0 == no mask, 1 == mask
+            src_mask = torch.zeros(
+                (src.shape[1], src.shape[1]),
+                device=src.device,
+                dtype=torch.bool
+            )
+
+        if left_context_chunks >= 0:
+            for i in range(src.shape[1]):
+                if chunk_size >= 0:
+                    current_chunk = (i // chunk_size) * chunk_size
+                    frame_remaining_context = max(0, current_chunk - left_context_chunks * chunk_size)
+                else:
+                    frame_remaining_context = 0
+
+                # end range is exclusive, so there is no off-by-one here
+                src_mask[i,:frame_remaining_context] = True
+
+        if chunk_size >= 0:
+            for i in range(src.shape[1]):
+                # if we have a chunk size of 8 then:
+                # for 0..7  -> mask 8..
+                # for 8..15 -> mask 16..
+                # etc.
+                visible_range = ((i // chunk_size) + 1) * chunk_size
+                src_mask[i, visible_range:] = True
+        else:
+            src_mask = None
 
         if self.causal:
             src_mask = get_lookahead_mask(src)
@@ -302,7 +332,7 @@ def decode(self, tgt, encoder_out, enc_len=None):
         )
         return prediction, multihead_attns[-1]
 
-    def encode(self, src, wav_len=None, pad_idx=0):
+    def encode(self, src, wav_len=None, pad_idx=0, chunk_size=None, left_context_chunks: int = -1):
         """
         Encoder forward pass
 
@@ -319,7 +349,7 @@ def encode(self, src, wav_len=None, pad_idx=0):
             src = src.reshape(bz, t, ch1 * ch2)
 
         (src_key_padding_mask, _, src_mask, _,) = self.make_masks(
-            src, None, wav_len, pad_idx=pad_idx
+            src, None, wav_len, pad_idx=pad_idx, chunk_size=chunk_size, left_context_chunks=left_context_chunks
         )
 
         src = self.custom_src_module(src)
@@ -333,8 +363,10 @@ def encode(self, src, wav_len=None, pad_idx=0):
 
         encoder_out, _ = self.encoder(
             src=src,
+            src_mask=src_mask,
             src_key_padding_mask=src_key_padding_mask,
             pos_embs=pos_embs_source,
+            chunk_size=chunk_size,
         )
         return encoder_out
 
@@ -373,7 +405,7 @@ def __init__(self, transformer, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.transformer = transformer
 
-    def forward(self, x, wav_lens=None, pad_idx=0):
+    def forward(self, x, wav_lens=None, pad_idx=0, chunk_size=-1, left_context_chunks=-1):
         """ Processes the input tensor x and returns an output tensor."""
-        x = self.transformer.encode(x, wav_lens, pad_idx)
+        x = self.transformer.encode(x, wav_lens, pad_idx, chunk_size=chunk_size, left_context_chunks=left_context_chunks)
         return x
diff --git a/speechbrain/nnet/attention.py b/speechbrain/nnet/attention.py
index 528522f673..64dd38f74f 100644
--- a/speechbrain/nnet/attention.py
+++ b/speechbrain/nnet/attention.py
@@ -624,6 +624,23 @@ def forward(
 
         attn_score = F.softmax(attn_score, dim=-1)
         attn_score = self.dropout_att(attn_score)
+
+        # it is possible for us to hit full NaN when using chunked training
+        # so reapply masks, except with 0.0 instead as we are after the softmax
+        # because -inf would output 0.0 regardless anyway
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                attn_score = attn_score.masked_fill(
+                    attn_mask, 0.0
+                )
+            else:
+                assert False, "oopsie need to reimplement that"
+
+        if key_padding_mask is not None:
+            attn_score = attn_score.masked_fill(
+                key_padding_mask.view(bsz, 1, 1, klen), 0.0,
+            )
+
         x = torch.matmul(
             attn_score, value.transpose(1, 2)
         )  # (batch, head, time1, d_k)

From b59accbb1a9073d5a71e8ae8a1889238ee0117d4 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 5 Jul 2023 11:26:41 +0200
Subject: [PATCH 02/83] DDP fix?

---
 speechbrain/lobes/models/transformer/TransformerASR.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 7e248acbf9..0dde3e03a7 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -152,9 +152,10 @@ def __init__(
             ),
             torch.nn.Dropout(dropout),
         )
-        self.custom_tgt_module = ModuleList(
-            NormalizedEmbedding(d_model, tgt_vocab)
-        )
+        if num_decoder_layers > 0:
+            self.custom_tgt_module = ModuleList(
+                NormalizedEmbedding(d_model, tgt_vocab)
+            )
 
         # reset parameters using xavier_normal_
         self._init_params()

From 9241781678b32aff8171808de9404ec039634423 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Sat, 8 Jul 2023 20:00:57 +0200
Subject: [PATCH 03/83] Batch of changes and things brought back

---
 .../transducer/hparams/conformer_transducer.yaml   |  3 ++-
 recipes/LibriSpeech/ASR/transducer/train.py        |  4 ++--
 speechbrain/lobes/models/transformer/Conformer.py  |  1 +
 speechbrain/nnet/attention.py                      | 14 ++++++++++----
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index 917cd95c39..7781c9885f 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -9,7 +9,7 @@
 # ############################################################################
 
 # Seed needs to be set at top of yaml, before objects with parameters are made
-seed: 3407
+seed: 1234
 __set_seed: !!python/object/apply:torch.manual_seed [!ref <seed>]
 output_folder: !ref results/conformer_transducer_large/<seed>
 output_wer_folder: !ref <output_folder>/
@@ -92,6 +92,7 @@ dynamic_left_context_max: 64
 train_dataloader_opts:
    batch_size: !ref <batch_size>
    num_workers: !ref <num_workers>
+   pin_memory: True
 
 valid_dataloader_opts:
    batch_size: !ref <batch_size_valid>
diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 9ad8bb699c..5f9b90e1d5 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -46,7 +46,7 @@
 class ASR(sb.Brain):
     def compute_forward(self, batch, stage):
         """Forward computations from the waveform batches to the output probabilities."""
-        batch = batch.to(self.device)
+        batch = batch.to(self.device, non_blocking=True)
         wavs, wav_lens = batch.sig
         tokens_with_bos, token_with_bos_lens = batch.tokens_bos
 
@@ -297,7 +297,7 @@ def on_evaluate_start(self, max_key=None, min_key=None):
         super().on_evaluate_start()
 
         ckpts = self.checkpointer.find_checkpoints(
-            max_key=max_key, min_key=min_key
+            max_key=max_key, min_key=min_key,
         )
         ckpt = sb.utils.checkpoints.average_checkpoints(
             ckpts, recoverable_name="model"
diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 7a7db13405..9662933054 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -145,6 +145,7 @@ def forward(self, x, mask=None, chunk_size=-1):
         # https://www.amazon.science/publications/dynamic-chunk-convolution-for-unified-streaming-and-non-streaming-conformer-asr
         # split the input into chunks of size `chunk_size`, but for each chunk
         # provide a left context for left chunk dependencies to be possible.
+
         if chunk_size >= 1:
             # chances are chunking+causal is unintended; i don't know where it
             # may make sense, but if it does to you, feel free to implement it.
diff --git a/speechbrain/nnet/attention.py b/speechbrain/nnet/attention.py
index 64dd38f74f..8430b04fab 100644
--- a/speechbrain/nnet/attention.py
+++ b/speechbrain/nnet/attention.py
@@ -591,17 +591,23 @@ def forward(
             query + self.pos_bias_v.view(1, 1, self.num_heads, self.head_dim)
         ).transpose(1, 2)
 
+        # TODO: cite https://asherliu.github.io/docs/sc21a.pdf
+        # for the scaling prior to the matrix multiplication 
+        # should read more of the paper though
+        # TODO: check if this causes any difference beyond precision
+        # (it does not seem like it does)
+
         # (batch, head, qlen, klen)
-        matrix_ac = torch.matmul(q_with_bias_u, key.permute(0, 2, 3, 1))
+        matrix_ac = torch.matmul(q_with_bias_u * self.scale, key.permute(0, 2, 3, 1))
         # (batch, num_heads, klen, 2*klen-1)
-        matrix_bd = torch.matmul(q_with_bias_v, p_k.permute(0, 2, 3, 1))
+        matrix_bd = torch.matmul(q_with_bias_v * self.scale, p_k.permute(0, 2, 3, 1))
         matrix_bd = self.rel_shift(matrix_bd)  # shifting trick
 
         # if klen != qlen:
         #   import ipdb
         #  ipdb.set_trace(
 
-        attn_score = (matrix_ac + matrix_bd) * self.scale
+        attn_score = (matrix_ac + matrix_bd) # already scaled above
 
         # compute attention probability
         if attn_mask is not None:
@@ -622,7 +628,7 @@ def forward(
                 key_padding_mask.view(bsz, 1, 1, klen), self.attn_fill_value,
             )
 
-        attn_score = F.softmax(attn_score, dim=-1)
+        attn_score = F.softmax(attn_score, dim=-1, dtype=torch.float32)
         attn_score = self.dropout_att(attn_score)
 
         # it is possible for us to hit full NaN when using chunked training

From a50454120c1d442cb0cba725bb52895e15a5b763 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 12 Jul 2023 14:43:15 +0200
Subject: [PATCH 04/83] Streaming fixes (successfully trains)

---
 .../hparams/conformer_transducer.yaml         |   4 +-
 recipes/LibriSpeech/ASR/transducer/train.py   |   8 +-
 .../models/transformer/TransformerASR.py      |   2 +-
 speechbrain/utils/streaming.py                | 158 ++++++++++++++++++
 4 files changed, 168 insertions(+), 4 deletions(-)
 create mode 100644 speechbrain/utils/streaming.py

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index 7781c9885f..1a9a6cc551 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -72,9 +72,9 @@ n_mels: 80
 win_length: 32
 
 # Streaming & dynamic chunk training options
-streaming: False  # controls all DCT & chunk size & left context mechanisms
+streaming: True  # controls all DCT & chunk size & left context mechanisms
 
-test_chunk_size: 8
+test_chunk_size: 16
 test_left_context_size: 64
 
 valid_chunk_size: -1
diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 5f9b90e1d5..820f648e9c 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -99,7 +99,13 @@ def compute_forward(self, batch, stage):
         feats = self.modules.normalize(feats, wav_lens, epoch=current_epoch)
 
         src = self.modules.CNN(feats)
-        x = self.modules.enc(src, wav_lens, pad_idx=self.hparams.pad_index)
+        x = self.modules.enc(
+            src,
+            wav_lens,
+            pad_idx=self.hparams.pad_index,
+            chunk_size=transformer_chunk_size,
+            left_context_chunks=left_context_chunks
+        )
         x = self.modules.proj_enc(x)
 
         e_in = self.modules.emb(tokens_with_bos)
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 0dde3e03a7..28ff042010 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -333,7 +333,7 @@ def decode(self, tgt, encoder_out, enc_len=None):
         )
         return prediction, multihead_attns[-1]
 
-    def encode(self, src, wav_len=None, pad_idx=0, chunk_size=None, left_context_chunks: int = -1):
+    def encode(self, src, wav_len=None, pad_idx=0, chunk_size=-1, left_context_chunks: int = -1):
         """
         Encoder forward pass
 
diff --git a/speechbrain/utils/streaming.py b/speechbrain/utils/streaming.py
new file mode 100644
index 0000000000..81257025ef
--- /dev/null
+++ b/speechbrain/utils/streaming.py
@@ -0,0 +1,158 @@
+"""Utilities to assist with designing and training streaming models.
+
+Authors
+* Sylvain de Langen 2023
+"""
+
+import math
+import torch
+from typing import Callable
+
+
+def chunkify_sequence(x, chunk_size):
+    chunks = []
+
+    for i in range(max(1, math.ceil(x.shape[1] / chunk_size))):
+        start = i * chunk_size
+        end = start + chunk_size
+        chunks.append(x[:,start:end,...])
+
+    return chunks
+
+
+def merge_chunks(chunks):
+    return torch.cat(chunks, 1)
+
+
+def chunked_wav_lens(chunks, wav_lens):
+    chunk_wav_lens = []
+
+    # consider 3 chunks: we have chunk_frac at 0.0, 0.333, 0.666
+    # for value 0.7:
+    # - the first two chunks are trivially 1.0.
+    # - the last chunk is (value - 0.666) / (1 / chunks)
+
+    for i in range(len(chunks)):
+        chunk_frac = i / len(chunks)
+        chunk_raw_len = (wav_lens - chunk_frac) * len(chunks)
+        chunk_raw_len = torch.clamp(chunk_raw_len, 0.0, 1.0)
+        chunk_wav_lens.append(chunk_raw_len)
+
+    return chunk_wav_lens
+
+
+def infer_dependency_matrix(
+    model: Callable,
+    seq_shape: tuple,
+    in_stride: int = 1
+):
+    """
+    Randomizes parts of the input sequence several times in order to detect
+    dependencies between input frames and output frames, aka whether a given
+    output frame depends on a given input frame.
+
+    This can prove useful to check whether a model behaves correctly in a
+    streaming context and does not contain accidental dependencies to future
+    frames that couldn't be known in a streaming scenario.
+
+    Note that this can get very computationally expensive for very long
+    sequences.
+
+    Furthermore, this expects inference to be fully deterministic, else false
+    dependencies may be found. This also means that the model must be in eval
+    mode, to inhibit things like dropout layers.
+
+    Arguments
+    ---------
+    model : Callable
+        Can be a model or a function (potentially emulating streaming
+        functionality). Does not require to be a trained model, random weights
+        should usually suffice.
+    seq_shape : tuple
+        The function tries inferring by randomizing parts of the input sequence
+        in order to detect unwanted dependencies.
+        The shape is expected to look like `[batch_size, seq_len, num_feats]`,
+        where `batch_size` may be `1`.
+    in_stride : int
+        Consider only N-th input, for when the input sequences are very long
+        (e.g. raw audio) and the output is shorter (subsampled, filters, etc.)
+
+    Returns
+    -------
+    dependencies : torch.BoolTensor
+        Matrix representing whether an output is dependent on an input; index
+        using `[in_frame_idx, out_frame_idx]`. `True` indicates a detected
+        dependency.
+    """
+    # TODO: document arguments
+
+    bs, seq_len, feat_len = seq_shape
+
+    base_seq = torch.rand(seq_shape)
+    with torch.no_grad():
+        base_out = model(base_seq)
+
+        if not model(base_seq).equal(base_out):
+            raise ValueError(
+                "Expected deterministic model, but inferring twice on the same "
+                "data yielded different results. Make sure that you use "
+                "`eval()` mode so that it does not include randomness."
+            )
+    out_len, _out_feat_len = base_out.shape[1:]
+
+    deps = torch.zeros(((seq_len + (in_stride - 1)) // in_stride, out_len), dtype=torch.bool)
+
+    for in_frame_idx in range(0, seq_len, in_stride):
+        test_seq = base_seq.clone()
+        test_seq[:,in_frame_idx,:] = torch.rand(bs, feat_len)
+
+        with torch.no_grad():
+            test_out = model(test_seq)
+
+        for out_frame_idx in range(out_len):
+            if not torch.allclose(
+                test_out[:,out_frame_idx,:],
+                base_out[:,out_frame_idx,:]
+            ):
+                deps[in_frame_idx // in_stride][out_frame_idx] = True
+
+    return deps
+
+def plot_dependency_matrix(deps, in_stride: int = 1):
+    """
+    Returns a matplotlib figure of a dependency matrix generated by
+    `infer_dependency_matrix`.
+
+    At a given point, a red square indicates that a given output frame (y-axis)
+    was to depend on a given input frame (x-axis).
+
+    For example, a fully red image means that all output frames were dependent
+    on all the history. This could be the case of a bidirectional RNN, or a
+    transformer model, for example.
+
+    Arguments
+    ---------
+    deps : torch.BoolTensor
+        Matrix returned by `infer_dependency_matrix` or one in a compatible
+        format.
+    """
+    import matplotlib.pyplot as plt
+    from matplotlib.colors import ListedColormap
+
+    cmap = ListedColormap(["white", "red"])
+
+    fig, ax = plt.subplots()
+    
+    ax.pcolormesh(
+        torch.permute(deps, (1, 0)),
+        cmap=cmap,
+        vmin=False,
+        vmax=True,
+        edgecolors="gray",
+        linewidth=0.5
+    )
+    ax.set_title("Dependency plot")
+    ax.set_xlabel("in")
+    ax.set_ylabel("out")
+    ax.set_aspect("equal")
+    return fig

From 008055bf2c173949783b26165f35a1e7a6ded8d3 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 18 Jul 2023 08:23:35 +0200
Subject: [PATCH 05/83] WIP streaming code

---
 .../lobes/models/transformer/Conformer.py     | 53 ++++++++++++++++++-
 .../models/transformer/TransformerASR.py      | 49 ++++++++++++++++-
 2 files changed, 100 insertions(+), 2 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 9662933054..cbabc93ea8 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -5,10 +5,11 @@
 * Samuele Cornell 2021
 """
 
+from dataclasses import dataclass
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from typing import Optional
+from typing import Optional, List
 import speechbrain as sb
 import math
 import warnings
@@ -24,6 +25,17 @@
 from speechbrain.nnet.activations import Swish
 
 
+@dataclass
+class ConformerEncoderLayerStreamingContext:
+    mha_left_context: Optional[torch.Tensor]
+    dcconv_left_context: Optional[torch.Tensor]
+
+
+@dataclass
+class ConformerEncoderStreamingContext:
+    layers: List[ConformerEncoderLayerStreamingContext]
+
+
 class ConvolutionModule(nn.Module):
     """This is an implementation of convolution module in Conformer.
 
@@ -388,6 +400,13 @@ def forward(
         return x, self_attn
 
 
+    def make_streaming_context(self):
+        return ConformerEncoderLayerStreamingContext(
+            mha_left_context=None,
+            dcconv_left_context=None
+        )
+
+
 class ConformerEncoder(nn.Module):
     """This class implements the Conformer encoder.
 
@@ -514,6 +533,38 @@ def forward(
 
         return output, attention_lst
 
+    def forward_streaming(
+        self,
+        src,
+        context: ConformerEncoderStreamingContext,
+        pos_embs: Optional[torch.Tensor] = None
+    ):
+        if self.attention_type == "RelPosMHAXL":
+            if pos_embs is None:
+                raise ValueError(
+                    "The chosen attention type for the Conformer is RelPosMHAXL. For this attention type, the positional embeddings are mandatory"
+                )
+
+        output = src
+        attention_lst = []
+        for i, enc_layer in enumerate(self.layers):
+            output, attention = enc_layer(
+                output,
+                pos_embs=pos_embs,
+                context=context.layers[i]
+            )
+            attention_lst.append(attention)
+        output = self.norm(output)
+
+        return output, attention_lst
+
+    def make_streaming_context(self):
+        return ConformerEncoderStreamingContext(
+            layers=[
+                layer.make_streaming_context() for layer in self.layers
+            ]
+        )
+
 
 class ConformerDecoderLayer(nn.Module):
     """This is an implementation of Conformer encoder layer.
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 28ff042010..2e63aecf3e 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -4,9 +4,10 @@
 * Jianyuan Zhong 2020
 """
 
+from dataclasses import dataclass
 import torch  # noqa 42
 from torch import nn
-from typing import Optional
+from typing import Any, Optional
 from speechbrain.nnet.linear import Linear
 from speechbrain.nnet.containers import ModuleList
 from speechbrain.lobes.models.transformer.Transformer import (
@@ -19,6 +20,20 @@
 from speechbrain.dataio.dataio import length_to_mask
 
 
+@dataclass
+class TransformerASRStreamingContext:
+    chunk_size: int
+    left_context_target_size: int
+    encoder: Any
+
+    @classmethod
+    def initial(cls, chunk_size, left_context_size):
+        return TransformerASRStreamingContext(
+            chunk_size=chunk_size,
+            left_context_target_size=left_context_size,
+        )
+
+
 class TransformerASR(TransformerInterface):
     """This is an implementation of transformer model for ASR.
 
@@ -371,6 +386,38 @@ def encode(self, src, wav_len=None, pad_idx=0, chunk_size=-1, left_context_chunk
         )
         return encoder_out
 
+    def encode_streaming(self, src, context: TransformerASRStreamingContext):
+        """
+        Streaming encoder forward pass
+        """
+        # TODO: docstring
+
+        if src.dim() == 4:
+            bz, t, ch1, ch2 = src.shape
+            src = src.reshape(bz, t, ch1 * ch2)
+
+        src = self.custom_src_module(src)
+        if self.attention_type == "RelPosMHAXL":
+            pos_embs_source = self.positional_encoding(src)
+
+        elif self.positional_encoding_type == "fixed_abs_sine":
+            src = src + self.positional_encoding(src)
+            pos_embs_source = None
+
+        encoder_out, _ = self.encoder.forward_streaming(
+            src=src,
+            pos_embs=pos_embs_source,
+            context=context.encoder
+        )
+        return encoder_out
+
+    def make_streaming_context(self, chunk_size, left_context_size, encoder_kwargs={}):
+        return TransformerASRStreamingContext(
+            chunk_size=chunk_size,
+            left_context_target_size=left_context_size,
+            encoder=self.encoder.make_streaming_context(**encoder_kwargs)
+        )
+
     def _init_params(self):
         for p in self.parameters():
             if p.dim() > 1:

From 4fc70f4cae0bee5684a68282a141659c688d8aa4 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 22 Aug 2023 11:47:26 +0200
Subject: [PATCH 06/83] WIP functional streaming code

---
 recipes/LibriSpeech/ASR/transducer/train.py   |  4 +-
 .../lobes/models/transformer/Conformer.py     | 52 ++++++++++++++++++-
 .../models/transformer/TransformerASR.py      | 41 ++++++++++-----
 3 files changed, 82 insertions(+), 15 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 820f648e9c..0b46d6134f 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -309,8 +309,8 @@ def on_evaluate_start(self, max_key=None, min_key=None):
             ckpts, recoverable_name="model"
         )
 
-        self.hparams.model.load_state_dict(ckpt, strict=True)
-        self.hparams.model.eval()
+        # self.hparams.model.load_state_dict(ckpt, strict=True)
+        # self.hparams.model.eval()
 
 
 def dataio_prepare(hparams):
diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index cbabc93ea8..216ee6eea8 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -399,6 +399,56 @@ def forward(
         x = self.norm2(x + 0.5 * self.ffn_module2(x))
         return x, self_attn
 
+    def streaming_forward(
+        self,
+        x,
+        context: ConformerEncoderLayerStreamingContext,
+        pos_embs: torch.Tensor = None,
+    ):
+        orig_len = x.shape[-2]
+        # print("pre ffn x:  ", x.shape)
+        # ffn module
+        x = x + 0.5 * self.ffn_module1(x)
+
+        # print("x:          ", x.shape)
+        if context.mha_left_context is not None:
+            x = torch.cat((context.mha_left_context, x), dim=1)
+
+        # print("cat(lc, x): ", x.shape)
+        context.mha_left_context = x[...,-32:,:]
+
+        # print("pos_embs:   ", pos_embs.shape)
+        # print("new lc:     ", context.mha_left_context.shape)
+        # print()
+
+        # muti-head attention module
+        skip = x
+        x = self.norm1(x)
+
+        x, self_attn = self.mha_layer(
+            x,
+            x,
+            x,
+            attn_mask=None,
+            key_padding_mask=None,
+            pos_embs=pos_embs,
+        )
+        x = x + skip
+        x = x[...,-orig_len:,:]
+
+        if context.dcconv_left_context is not None:
+            x = torch.cat((context.dcconv_left_context, x), dim=1)
+
+        context.dcconv_left_context = x[...,-self.convolution_module.padding:,:]
+
+        # convolution module
+        x = x + self.convolution_module(x)
+
+        x = x[...,-orig_len:,:]
+
+        # ffn module
+        x = self.norm2(x + 0.5 * self.ffn_module2(x))
+        return x, self_attn
 
     def make_streaming_context(self):
         return ConformerEncoderLayerStreamingContext(
@@ -548,7 +598,7 @@ def forward_streaming(
         output = src
         attention_lst = []
         for i, enc_layer in enumerate(self.layers):
-            output, attention = enc_layer(
+            output, attention = enc_layer.streaming_forward(
                 output,
                 pos_embs=pos_embs,
                 context=context.layers[i]
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 2e63aecf3e..07b88ad230 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -24,14 +24,7 @@
 class TransformerASRStreamingContext:
     chunk_size: int
     left_context_target_size: int
-    encoder: Any
-
-    @classmethod
-    def initial(cls, chunk_size, left_context_size):
-        return TransformerASRStreamingContext(
-            chunk_size=chunk_size,
-            left_context_target_size=left_context_size,
-        )
+    encoder_context: Any
 
 
 class TransformerASR(TransformerInterface):
@@ -396,18 +389,42 @@ def encode_streaming(self, src, context: TransformerASRStreamingContext):
             bz, t, ch1, ch2 = src.shape
             src = src.reshape(bz, t, ch1 * ch2)
 
+        # HACK: our problem here is that the positional_encoding is computed
+        # against the size of our source tensor, but we only know how many left
+        # context frames we're injecting to the encoder within the encoder
+        # context.
+        # so this workaround does just that.
+        #
+        # i'm not sure how this would be best refactored, but an option would be
+        # to let the encoder get the pos embedding itself and have a way to
+        # cache it.
+        #
+        # additionally, positional encoding functions take in a whole source
+        # tensor just to get its attributes (size, device, type) but this is
+        # sort of silly for the embeddings that don't need one.
+        # so we craft a dummy empty (uninitialized) tensor to help...
+        known_left_context = context.encoder_context.layers[0].mha_left_context
+        if known_left_context is None:
+            # print(f"no lc known: using {src.shape}")
+            pos_encoding_dummy = src
+        else:
+            target_shape = list(src.shape)
+            # print(f"computing posemb shape: from {target_shape} with lc {known_left_context.shape}")
+            target_shape[-2] += known_left_context.shape[-2]
+            pos_encoding_dummy = torch.empty(size=target_shape).to(src)
+
         src = self.custom_src_module(src)
         if self.attention_type == "RelPosMHAXL":
-            pos_embs_source = self.positional_encoding(src)
+            pos_embs_source = self.positional_encoding(pos_encoding_dummy)
 
         elif self.positional_encoding_type == "fixed_abs_sine":
-            src = src + self.positional_encoding(src)
+            src = src + self.positional_encoding(pos_encoding_dummy)
             pos_embs_source = None
 
         encoder_out, _ = self.encoder.forward_streaming(
             src=src,
             pos_embs=pos_embs_source,
-            context=context.encoder
+            context=context.encoder_context
         )
         return encoder_out
 
@@ -415,7 +432,7 @@ def make_streaming_context(self, chunk_size, left_context_size, encoder_kwargs={
         return TransformerASRStreamingContext(
             chunk_size=chunk_size,
             left_context_target_size=left_context_size,
-            encoder=self.encoder.make_streaming_context(**encoder_kwargs)
+            encoder_context=self.encoder.make_streaming_context(**encoder_kwargs)
         )
 
     def _init_params(self):

From 4a7e95f891c5333fda2af526155ddca1a05ce244 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 22 Aug 2023 16:06:38 +0200
Subject: [PATCH 07/83] Fix left context

---
 speechbrain/lobes/models/transformer/Conformer.py  | 14 ++++++++++----
 .../lobes/models/transformer/TransformerASR.py     |  6 +++++-
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 216ee6eea8..647d713c91 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -27,6 +27,7 @@
 
 @dataclass
 class ConformerEncoderLayerStreamingContext:
+    mha_left_context_size: int
     mha_left_context: Optional[torch.Tensor]
     dcconv_left_context: Optional[torch.Tensor]
 
@@ -415,7 +416,8 @@ def streaming_forward(
             x = torch.cat((context.mha_left_context, x), dim=1)
 
         # print("cat(lc, x): ", x.shape)
-        context.mha_left_context = x[...,-32:,:]
+        if context.mha_left_context_size > 0:
+            context.mha_left_context = x[...,-context.mha_left_context_size:,:]
 
         # print("pos_embs:   ", pos_embs.shape)
         # print("new lc:     ", context.mha_left_context.shape)
@@ -450,8 +452,9 @@ def streaming_forward(
         x = self.norm2(x + 0.5 * self.ffn_module2(x))
         return x, self_attn
 
-    def make_streaming_context(self):
+    def make_streaming_context(self, mha_left_context_size: int):
         return ConformerEncoderLayerStreamingContext(
+            mha_left_context_size=mha_left_context_size,
             mha_left_context=None,
             dcconv_left_context=None
         )
@@ -608,10 +611,13 @@ def forward_streaming(
 
         return output, attention_lst
 
-    def make_streaming_context(self):
+    def make_streaming_context(self, mha_left_context_size: int):
         return ConformerEncoderStreamingContext(
             layers=[
-                layer.make_streaming_context() for layer in self.layers
+                layer.make_streaming_context(
+                    mha_left_context_size=mha_left_context_size
+                )
+                for layer in self.layers
             ]
         )
 
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 07b88ad230..50108a36ab 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -432,7 +432,11 @@ def make_streaming_context(self, chunk_size, left_context_size, encoder_kwargs={
         return TransformerASRStreamingContext(
             chunk_size=chunk_size,
             left_context_target_size=left_context_size,
-            encoder_context=self.encoder.make_streaming_context(**encoder_kwargs)
+            encoder_context=self.encoder.make_streaming_context(
+                # FIXME: bad naming, not all encoders might use mha etc
+                mha_left_context_size=left_context_size,
+                **encoder_kwargs
+            )
         )
 
     def _init_params(self):

From 83616a1b6766699a7a4c35a6129b2321b40fbb48 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 23 Aug 2023 09:54:19 +0200
Subject: [PATCH 08/83] Fix formatting

---
 recipes/LibriSpeech/ASR/transducer/train.py   | 13 ++--
 .../lobes/models/transformer/Conformer.py     | 65 ++++++++++---------
 .../models/transformer/TransformerASR.py      | 58 +++++++++++++----
 speechbrain/nnet/attention.py                 | 16 +++--
 speechbrain/utils/streaming.py                | 20 +++---
 5 files changed, 105 insertions(+), 67 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 0b46d6134f..ec86736f48 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -74,19 +74,22 @@ def compute_forward(self, batch, stage):
         left_context_chunks = -1
         if self.hparams.streaming:
             if stage == sb.Stage.TRAIN:
-                if torch.rand((1, )).item() < self.hparams.dynamic_chunk_thresh:
+                if torch.rand((1,)).item() < self.hparams.dynamic_chunk_thresh:
                     transformer_chunk_size = torch.randint(
                         self.hparams.dynamic_chunk_min,
                         self.hparams.dynamic_chunk_max + 1,
-                        (1, )
+                        (1,),
                     ).item()
                     # print("tfx chunk size", transformer_chunk_size)
 
-                if torch.rand((1, )).item() < self.hparams.dynamic_left_context_thresh:
+                if (
+                    torch.rand((1,)).item()
+                    < self.hparams.dynamic_left_context_thresh
+                ):
                     left_context_chunks = torch.randint(
                         self.hparams.dynamic_left_context_min,
                         self.hparams.dynamic_left_context_max + 1,
-                        (1, )
+                        (1,),
                     ).item()
             elif stage == sb.Stage.TEST:
                 transformer_chunk_size = self.hparams.test_chunk_size
@@ -104,7 +107,7 @@ def compute_forward(self, batch, stage):
             wav_lens,
             pad_idx=self.hparams.pad_index,
             chunk_size=transformer_chunk_size,
-            left_context_chunks=left_context_chunks
+            left_context_chunks=left_context_chunks,
         )
         x = self.modules.proj_enc(x)
 
diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 647d713c91..8ddc9fa811 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -179,10 +179,7 @@ def forward(self, x, mask=None, chunk_size=-1):
             # compute the left context that can and should be added, for each
             # chunk. for the first few chunks, we will need to add extra padding
             applied_left_context = [
-                min(
-                    chunk_left_context,
-                    i * chunk_size,
-                )
+                min(chunk_left_context, i * chunk_size,)
                 for i in range(chunk_count)
             ]
 
@@ -190,23 +187,32 @@ def forward(self, x, mask=None, chunk_size=-1):
             # the left context effectively becomes "left padding", we do not
             # want to keep any convolution results centered on the left context
             out = [
-                x[:,i * chunk_size - applied_left_context[i]:(i + 1) * chunk_size,...]
+                x[
+                    :,
+                    i * chunk_size
+                    - applied_left_context[i] : (i + 1) * chunk_size,
+                    ...,
+                ]
                 for i in range(chunk_count)
             ]
 
             # TODO: experiment around reflect padding, which is difficult
             # because small chunks have too little time steps to reflect from
             out = [
-                F.pad(out[i], (
-                    # channel dims, we do not to pad these
-                    0,
-                    0,
-                    # add missing left 0-padding if we lacked left context
-                    chunk_left_context - applied_left_context[i],
-                    # add missing right 0-padding as we disable default padding
-                    # also add missing frames of the rightmost chunk
-                    self.padding + (final_right_padding if i == len(out) - 1 else 0)
-                ))
+                F.pad(
+                    out[i],
+                    (
+                        # channel dims, we do not to pad these
+                        0,
+                        0,
+                        # add missing left 0-padding if we lacked left context
+                        chunk_left_context - applied_left_context[i],
+                        # add missing right 0-padding as we disable default padding
+                        # also add missing frames of the rightmost chunk
+                        self.padding
+                        + (final_right_padding if i == len(out) - 1 else 0),
+                    ),
+                )
                 for i in range(len(out))
             ]
 
@@ -230,7 +236,7 @@ def forward(self, x, mask=None, chunk_size=-1):
 
             # -> [batch_size, time_steps, out_channels]
             if final_right_padding > 0:
-                out = out[:,:-final_right_padding,:]
+                out = out[:, :-final_right_padding, :]
         else:
             out = self._do_conv(x, inhibit_padding=False)
 
@@ -417,7 +423,9 @@ def streaming_forward(
 
         # print("cat(lc, x): ", x.shape)
         if context.mha_left_context_size > 0:
-            context.mha_left_context = x[...,-context.mha_left_context_size:,:]
+            context.mha_left_context = x[
+                ..., -context.mha_left_context_size :, :
+            ]
 
         # print("pos_embs:   ", pos_embs.shape)
         # print("new lc:     ", context.mha_left_context.shape)
@@ -428,25 +436,22 @@ def streaming_forward(
         x = self.norm1(x)
 
         x, self_attn = self.mha_layer(
-            x,
-            x,
-            x,
-            attn_mask=None,
-            key_padding_mask=None,
-            pos_embs=pos_embs,
+            x, x, x, attn_mask=None, key_padding_mask=None, pos_embs=pos_embs,
         )
         x = x + skip
-        x = x[...,-orig_len:,:]
+        x = x[..., -orig_len:, :]
 
         if context.dcconv_left_context is not None:
             x = torch.cat((context.dcconv_left_context, x), dim=1)
 
-        context.dcconv_left_context = x[...,-self.convolution_module.padding:,:]
+        context.dcconv_left_context = x[
+            ..., -self.convolution_module.padding :, :
+        ]
 
         # convolution module
         x = x + self.convolution_module(x)
 
-        x = x[...,-orig_len:,:]
+        x = x[..., -orig_len:, :]
 
         # ffn module
         x = self.norm2(x + 0.5 * self.ffn_module2(x))
@@ -456,7 +461,7 @@ def make_streaming_context(self, mha_left_context_size: int):
         return ConformerEncoderLayerStreamingContext(
             mha_left_context_size=mha_left_context_size,
             mha_left_context=None,
-            dcconv_left_context=None
+            dcconv_left_context=None,
         )
 
 
@@ -590,7 +595,7 @@ def forward_streaming(
         self,
         src,
         context: ConformerEncoderStreamingContext,
-        pos_embs: Optional[torch.Tensor] = None
+        pos_embs: Optional[torch.Tensor] = None,
     ):
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:
@@ -602,9 +607,7 @@ def forward_streaming(
         attention_lst = []
         for i, enc_layer in enumerate(self.layers):
             output, attention = enc_layer.streaming_forward(
-                output,
-                pos_embs=pos_embs,
-                context=context.layers[i]
+                output, pos_embs=pos_embs, context=context.layers[i]
             )
             attention_lst.append(attention)
         output = self.norm(output)
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 50108a36ab..b8716772ad 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -239,7 +239,15 @@ def forward(self, src, tgt, wav_len=None, pad_idx=0):
 
         return encoder_out, decoder_out
 
-    def make_masks(self, src, tgt=None, wav_len=None, pad_idx=0, chunk_size: int = -1, left_context_chunks: int = -1):
+    def make_masks(
+        self,
+        src,
+        tgt=None,
+        wav_len=None,
+        pad_idx=0,
+        chunk_size: int = -1,
+        left_context_chunks: int = -1,
+    ):
         """This method generates the masks for training the transformer model.
 
         Arguments
@@ -262,19 +270,21 @@ def make_masks(self, src, tgt=None, wav_len=None, pad_idx=0, chunk_size: int = -
             src_mask = torch.zeros(
                 (src.shape[1], src.shape[1]),
                 device=src.device,
-                dtype=torch.bool
+                dtype=torch.bool,
             )
 
         if left_context_chunks >= 0:
             for i in range(src.shape[1]):
                 if chunk_size >= 0:
                     current_chunk = (i // chunk_size) * chunk_size
-                    frame_remaining_context = max(0, current_chunk - left_context_chunks * chunk_size)
+                    frame_remaining_context = max(
+                        0, current_chunk - left_context_chunks * chunk_size
+                    )
                 else:
                     frame_remaining_context = 0
 
                 # end range is exclusive, so there is no off-by-one here
-                src_mask[i,:frame_remaining_context] = True
+                src_mask[i, :frame_remaining_context] = True
 
         if chunk_size >= 0:
             for i in range(src.shape[1]):
@@ -341,7 +351,14 @@ def decode(self, tgt, encoder_out, enc_len=None):
         )
         return prediction, multihead_attns[-1]
 
-    def encode(self, src, wav_len=None, pad_idx=0, chunk_size=-1, left_context_chunks: int = -1):
+    def encode(
+        self,
+        src,
+        wav_len=None,
+        pad_idx=0,
+        chunk_size=-1,
+        left_context_chunks: int = -1,
+    ):
         """
         Encoder forward pass
 
@@ -358,7 +375,12 @@ def encode(self, src, wav_len=None, pad_idx=0, chunk_size=-1, left_context_chunk
             src = src.reshape(bz, t, ch1 * ch2)
 
         (src_key_padding_mask, _, src_mask, _,) = self.make_masks(
-            src, None, wav_len, pad_idx=pad_idx, chunk_size=chunk_size, left_context_chunks=left_context_chunks
+            src,
+            None,
+            wav_len,
+            pad_idx=pad_idx,
+            chunk_size=chunk_size,
+            left_context_chunks=left_context_chunks,
         )
 
         src = self.custom_src_module(src)
@@ -422,21 +444,21 @@ def encode_streaming(self, src, context: TransformerASRStreamingContext):
             pos_embs_source = None
 
         encoder_out, _ = self.encoder.forward_streaming(
-            src=src,
-            pos_embs=pos_embs_source,
-            context=context.encoder_context
+            src=src, pos_embs=pos_embs_source, context=context.encoder_context
         )
         return encoder_out
 
-    def make_streaming_context(self, chunk_size, left_context_size, encoder_kwargs={}):
+    def make_streaming_context(
+        self, chunk_size, left_context_size, encoder_kwargs={}
+    ):
         return TransformerASRStreamingContext(
             chunk_size=chunk_size,
             left_context_target_size=left_context_size,
             encoder_context=self.encoder.make_streaming_context(
                 # FIXME: bad naming, not all encoders might use mha etc
                 mha_left_context_size=left_context_size,
-                **encoder_kwargs
-            )
+                **encoder_kwargs,
+            ),
         )
 
     def _init_params(self):
@@ -474,7 +496,15 @@ def __init__(self, transformer, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.transformer = transformer
 
-    def forward(self, x, wav_lens=None, pad_idx=0, chunk_size=-1, left_context_chunks=-1):
+    def forward(
+        self, x, wav_lens=None, pad_idx=0, chunk_size=-1, left_context_chunks=-1
+    ):
         """ Processes the input tensor x and returns an output tensor."""
-        x = self.transformer.encode(x, wav_lens, pad_idx, chunk_size=chunk_size, left_context_chunks=left_context_chunks)
+        x = self.transformer.encode(
+            x,
+            wav_lens,
+            pad_idx,
+            chunk_size=chunk_size,
+            left_context_chunks=left_context_chunks,
+        )
         return x
diff --git a/speechbrain/nnet/attention.py b/speechbrain/nnet/attention.py
index 8430b04fab..74d1305a6f 100644
--- a/speechbrain/nnet/attention.py
+++ b/speechbrain/nnet/attention.py
@@ -592,22 +592,26 @@ def forward(
         ).transpose(1, 2)
 
         # TODO: cite https://asherliu.github.io/docs/sc21a.pdf
-        # for the scaling prior to the matrix multiplication 
+        # for the scaling prior to the matrix multiplication
         # should read more of the paper though
         # TODO: check if this causes any difference beyond precision
         # (it does not seem like it does)
 
         # (batch, head, qlen, klen)
-        matrix_ac = torch.matmul(q_with_bias_u * self.scale, key.permute(0, 2, 3, 1))
+        matrix_ac = torch.matmul(
+            q_with_bias_u * self.scale, key.permute(0, 2, 3, 1)
+        )
         # (batch, num_heads, klen, 2*klen-1)
-        matrix_bd = torch.matmul(q_with_bias_v * self.scale, p_k.permute(0, 2, 3, 1))
+        matrix_bd = torch.matmul(
+            q_with_bias_v * self.scale, p_k.permute(0, 2, 3, 1)
+        )
         matrix_bd = self.rel_shift(matrix_bd)  # shifting trick
 
         # if klen != qlen:
         #   import ipdb
         #  ipdb.set_trace(
 
-        attn_score = (matrix_ac + matrix_bd) # already scaled above
+        attn_score = matrix_ac + matrix_bd  # already scaled above
 
         # compute attention probability
         if attn_mask is not None:
@@ -636,9 +640,7 @@ def forward(
         # because -inf would output 0.0 regardless anyway
         if attn_mask is not None:
             if attn_mask.dtype == torch.bool:
-                attn_score = attn_score.masked_fill(
-                    attn_mask, 0.0
-                )
+                attn_score = attn_score.masked_fill(attn_mask, 0.0)
             else:
                 assert False, "oopsie need to reimplement that"
 
diff --git a/speechbrain/utils/streaming.py b/speechbrain/utils/streaming.py
index 81257025ef..12fd76fd66 100644
--- a/speechbrain/utils/streaming.py
+++ b/speechbrain/utils/streaming.py
@@ -15,7 +15,7 @@ def chunkify_sequence(x, chunk_size):
     for i in range(max(1, math.ceil(x.shape[1] / chunk_size))):
         start = i * chunk_size
         end = start + chunk_size
-        chunks.append(x[:,start:end,...])
+        chunks.append(x[:, start:end, ...])
 
     return chunks
 
@@ -42,9 +42,7 @@ def chunked_wav_lens(chunks, wav_lens):
 
 
 def infer_dependency_matrix(
-    model: Callable,
-    seq_shape: tuple,
-    in_stride: int = 1
+    model: Callable, seq_shape: tuple, in_stride: int = 1
 ):
     """
     Randomizes parts of the input sequence several times in order to detect
@@ -100,24 +98,26 @@ def infer_dependency_matrix(
             )
     out_len, _out_feat_len = base_out.shape[1:]
 
-    deps = torch.zeros(((seq_len + (in_stride - 1)) // in_stride, out_len), dtype=torch.bool)
+    deps = torch.zeros(
+        ((seq_len + (in_stride - 1)) // in_stride, out_len), dtype=torch.bool
+    )
 
     for in_frame_idx in range(0, seq_len, in_stride):
         test_seq = base_seq.clone()
-        test_seq[:,in_frame_idx,:] = torch.rand(bs, feat_len)
+        test_seq[:, in_frame_idx, :] = torch.rand(bs, feat_len)
 
         with torch.no_grad():
             test_out = model(test_seq)
 
         for out_frame_idx in range(out_len):
             if not torch.allclose(
-                test_out[:,out_frame_idx,:],
-                base_out[:,out_frame_idx,:]
+                test_out[:, out_frame_idx, :], base_out[:, out_frame_idx, :]
             ):
                 deps[in_frame_idx // in_stride][out_frame_idx] = True
 
     return deps
 
+
 def plot_dependency_matrix(deps, in_stride: int = 1):
     """
     Returns a matplotlib figure of a dependency matrix generated by
@@ -142,14 +142,14 @@ def plot_dependency_matrix(deps, in_stride: int = 1):
     cmap = ListedColormap(["white", "red"])
 
     fig, ax = plt.subplots()
-    
+
     ax.pcolormesh(
         torch.permute(deps, (1, 0)),
         cmap=cmap,
         vmin=False,
         vmax=True,
         edgecolors="gray",
-        linewidth=0.5
+        linewidth=0.5,
     )
     ax.set_title("Dependency plot")
     ax.set_xlabel("in")

From 9b3a0d3643631590c43c561e8a623679e5172487 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 23 Aug 2023 11:27:51 +0200
Subject: [PATCH 09/83] Cleanups and docs in streaming utils

---
 speechbrain/utils/streaming.py | 115 +++++++++++++++++++++++++++------
 1 file changed, 96 insertions(+), 19 deletions(-)

diff --git a/speechbrain/utils/streaming.py b/speechbrain/utils/streaming.py
index 12fd76fd66..15911db451 100644
--- a/speechbrain/utils/streaming.py
+++ b/speechbrain/utils/streaming.py
@@ -6,38 +6,115 @@
 
 import math
 import torch
-from typing import Callable
+from typing import Callable, List
 
 
-def chunkify_sequence(x, chunk_size):
-    chunks = []
+def split_fixed_chunks(
+    x: torch.Tensor, chunk_size: int, dim: int = -1
+) -> List[torch.Tensor]:
+    """Split an input tensor `x` into a list of chunk tensors of size
+    `chunk_size` alongside dimension `dim`.
+    Useful for splitting up sequences with chunks of fixed sizes.
 
-    for i in range(max(1, math.ceil(x.shape[1] / chunk_size))):
-        start = i * chunk_size
-        end = start + chunk_size
-        chunks.append(x[:, start:end, ...])
+    If dimension `dim` cannot be evenly split by `chunk_size`, then the last
+    chunk will be smaller than `chunk_size`.
 
-    return chunks
+    Arguments
+    ---------
+    x : torch.Tensor
+        The tensor to split into chunks, typically a sequence or audio signal.
+
+    chunk_size : int
+        The size of each chunk, i.e. the max size of each chunk on dimension
+        `dim`.
+
+    dim : int
+        Dimension to split alongside of, typically the time dimension.
+
+    Returns
+    -------
+    List[torch.Tensor]
+        A chunk list of tensors, see description and example.
+        Guarantees `.size(dim) <= chunk_size`.
+
+    Example
+    -------
+    >>> import torch
+    >>> from speechbrain.utils.streaming import split_fixed_chunks
+    >>> x = torch.zeros((16, 10000, 80))
+    >>> chunks = split_fixed_chunks(x, 128, dim=1)
+    >>> len(chunks)
+    79
+    >>> chunks[0].shape
+    torch.Size([16, 128, 80])
+    >>> chunks[-1].shape
+    torch.Size([16, 16, 80])
+    """
+
+    num_chunks = math.ceil(x.size(dim) / chunk_size)
+    split_at_indices = [(i + 1) * chunk_size for i in range(num_chunks - 1)]
+    return torch.tensor_split(x, split_at_indices, dim=1)
 
 
-def merge_chunks(chunks):
-    return torch.cat(chunks, 1)
+def split_wav_lens(
+    chunk_lens: List[int], wav_lens: torch.Tensor
+) -> List[torch.Tensor]:
+    """Converts a single `wav_lens` tensor into a list of `chunk_count` tensors,
+    typically useful when chunking signals with `split_fixed_chunks`.
 
+    `wav_lens` represents the relative length of each audio within a batch,
+    which is typically used for masking. This function computes the relative
+    length at chunk level.
+
+    Arguments
+    ---------
+    chunk_lens : List[int]
+        Length of the sequence of every chunk. For example, if `chunks` was
+        returned from `split_fixed_chunks(x, chunk_size, dim=1)`, then this
+        should be `[chk.size(1) for chk in chunks]`.
+
+    wav_lens : torch.Tensor
+        Relative lengths of audio within a batch. For example, for an input
+        signal of 100 frames and a batch of 3 elements, `(1.0, 0.5, 0.25)`
+        would mean the batch holds audio of 100 frames, 50 frames and 25 frames
+        respectively.
+
+    Returns
+    -------
+    List[torch.Tensor]
+        A list of chunked wav_lens, see description and example.
+
+    Example
+    -------
+    >>> import torch
+    >>> from speechbrain.utils.streaming import split_wav_lens, split_fixed_chunks
+    >>> x = torch.zeros((3, 20, 80))
+    >>> chunks = split_fixed_chunks(x, 8, dim=1)
+    >>> len(chunks)
+    3
+    >>> # 20 frames, 13 frames, 17 frames
+    >>> wav_lens = torch.tensor([1.0, 0.65, 0.85])
+    >>> chunked_wav_lens = split_wav_lens([c.size(1) for c in chunks], wav_lens)
+    >>> chunked_wav_lens
+    [tensor([1., 1., 1.]),
+     tensor([1.0000, 0.6250, 1.0000]),
+     tensor([1.0000, 0.0000, 0.2500])]
+    >>> # wav 1 covers 62.5% (5/8) of the second chunk's frames
+    """
 
-def chunked_wav_lens(chunks, wav_lens):
     chunk_wav_lens = []
 
-    # consider 3 chunks: we have chunk_frac at 0.0, 0.333, 0.666
-    # for value 0.7:
-    # - the first two chunks are trivially 1.0.
-    # - the last chunk is (value - 0.666) / (1 / chunks)
+    seq_size = sum(chunk_lens)
+    wav_lens_frames = wav_lens * seq_size
 
-    for i in range(len(chunks)):
-        chunk_frac = i / len(chunks)
-        chunk_raw_len = (wav_lens - chunk_frac) * len(chunks)
+    chunk_start_frame = 0
+    for chunk_len in chunk_lens:
+        chunk_raw_len = (wav_lens_frames - chunk_start_frame) / chunk_len
         chunk_raw_len = torch.clamp(chunk_raw_len, 0.0, 1.0)
         chunk_wav_lens.append(chunk_raw_len)
 
+        chunk_start_frame += chunk_len
+
     return chunk_wav_lens
 
 
@@ -118,7 +195,7 @@ def infer_dependency_matrix(
     return deps
 
 
-def plot_dependency_matrix(deps, in_stride: int = 1):
+def plot_dependency_matrix(deps):
     """
     Returns a matplotlib figure of a dependency matrix generated by
     `infer_dependency_matrix`.

From 4ee692f7a543a94a3e3ba4cdab7a8a8f8103602b Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 23 Aug 2023 12:07:55 +0200
Subject: [PATCH 10/83] Better comment hparams, change seed back to orig,
 improve naming

---
 .../ASR/transducer/hparams/conformer_transducer.yaml     | 9 ++++++---
 recipes/LibriSpeech/ASR/transducer/train.py              | 4 ++--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index 1a9a6cc551..1b28205ddc 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -9,7 +9,7 @@
 # ############################################################################
 
 # Seed needs to be set at top of yaml, before objects with parameters are made
-seed: 1234
+seed: 3407
 __set_seed: !!python/object/apply:torch.manual_seed [!ref <seed>]
 output_folder: !ref results/conformer_transducer_large/<seed>
 output_wer_folder: !ref <output_folder>/
@@ -72,6 +72,9 @@ n_mels: 80
 win_length: 32
 
 # Streaming & dynamic chunk training options
+# At least for the current architecture on LibriSpeech, we found out that
+# non-streaming accuracy is very similar between `streaming: True` and
+# `streaming: False`.
 streaming: True  # controls all DCT & chunk size & left context mechanisms
 
 test_chunk_size: 16
@@ -80,11 +83,11 @@ test_left_context_size: 64
 valid_chunk_size: -1
 valid_left_context_size: -1
 
-dynamic_chunk_thresh: 0.6
+dynamic_chunk_prob: 0.6
 dynamic_chunk_min: 8
 dynamic_chunk_max: 32
 
-dynamic_left_context_thresh: 0.75 # TODO: rename all to prob, makes more sense
+dynamic_left_context_prob: 0.75 # TODO: rename all to prob, makes more sense
 dynamic_left_context_min: 16
 dynamic_left_context_max: 64
 
diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index ec86736f48..84eb7f500a 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -74,7 +74,7 @@ def compute_forward(self, batch, stage):
         left_context_chunks = -1
         if self.hparams.streaming:
             if stage == sb.Stage.TRAIN:
-                if torch.rand((1,)).item() < self.hparams.dynamic_chunk_thresh:
+                if torch.rand((1,)).item() < self.hparams.dynamic_chunk_prob:
                     transformer_chunk_size = torch.randint(
                         self.hparams.dynamic_chunk_min,
                         self.hparams.dynamic_chunk_max + 1,
@@ -84,7 +84,7 @@ def compute_forward(self, batch, stage):
 
                 if (
                     torch.rand((1,)).item()
-                    < self.hparams.dynamic_left_context_thresh
+                    < self.hparams.dynamic_left_context_prob
                 ):
                     left_context_chunks = torch.randint(
                         self.hparams.dynamic_left_context_min,

From b459187725a51b11e05a45b498b8c8fa295a3ff2 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 23 Aug 2023 12:10:13 +0200
Subject: [PATCH 11/83] uncomment averaging stuff; it was some ipython issue

---
 recipes/LibriSpeech/ASR/transducer/train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 84eb7f500a..f677bd9905 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -312,8 +312,8 @@ def on_evaluate_start(self, max_key=None, min_key=None):
             ckpts, recoverable_name="model"
         )
 
-        # self.hparams.model.load_state_dict(ckpt, strict=True)
-        # self.hparams.model.eval()
+        self.hparams.model.load_state_dict(ckpt, strict=True)
+        self.hparams.model.eval()
 
 
 def dataio_prepare(hparams):

From fa5edea12fdeadfa6514af65674b1f8c30510898 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 23 Aug 2023 12:12:02 +0200
Subject: [PATCH 12/83] Remove pin_memory as it was not beneficial

---
 .../ASR/transducer/hparams/conformer_transducer.yaml            | 1 -
 recipes/LibriSpeech/ASR/transducer/train.py                     | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index 1b28205ddc..a4957844a3 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -95,7 +95,6 @@ dynamic_left_context_max: 64
 train_dataloader_opts:
    batch_size: !ref <batch_size>
    num_workers: !ref <num_workers>
-   pin_memory: True
 
 valid_dataloader_opts:
    batch_size: !ref <batch_size_valid>
diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index f677bd9905..7f11fb8166 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -46,7 +46,7 @@
 class ASR(sb.Brain):
     def compute_forward(self, batch, stage):
         """Forward computations from the waveform batches to the output probabilities."""
-        batch = batch.to(self.device, non_blocking=True)
+        batch = batch.to(self.device)
         wavs, wav_lens = batch.sig
         tokens_with_bos, token_with_bos_lens = batch.tokens_bos
 

From ed367766fdf3677c69ee584f5e807df74cecc9b9 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 23 Aug 2023 15:01:43 +0200
Subject: [PATCH 13/83] More cleanups, comments on context stuff

---
 recipes/LibriSpeech/ASR/transducer/train.py   | 15 ++++-
 .../lobes/models/transformer/Conformer.py     | 62 ++++++++++++++-----
 .../models/transformer/TransformerASR.py      | 27 ++++++--
 3 files changed, 83 insertions(+), 21 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 7f11fb8166..2b5eaae8e3 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -70,18 +70,31 @@ def compute_forward(self, batch, stage):
 
         current_epoch = self.hparams.epoch_counter.current
 
+        # Default to infinite context visibility
         transformer_chunk_size = -1
         left_context_chunks = -1
         if self.hparams.streaming:
+            # TODO: while this is fairly small logic, it may make sense to
+            # extract it to its own class orchestrating dynamic chunk training,
+            # partly because explantions are beneficial
             if stage == sb.Stage.TRAIN:
+                # When training for streaming, for each batch, we have a
+                # `dynamic_chunk_prob` probability of sampling a chunk size
+                # between `dynamic_chunk_min` and `_max`, otherwise output
+                # frames can see anywhere in the future.
+                # NOTE: We use torch random to be bound to the experiment seed.
                 if torch.rand((1,)).item() < self.hparams.dynamic_chunk_prob:
                     transformer_chunk_size = torch.randint(
                         self.hparams.dynamic_chunk_min,
                         self.hparams.dynamic_chunk_max + 1,
                         (1,),
                     ).item()
-                    # print("tfx chunk size", transformer_chunk_size)
 
+                # We have a `dynamic_left_context_prob` probability of sampling
+                # a left context size between `dynamic_left_context_min` and
+                # `_max`, otherwise output frames can see anywhere in the past.
+                # Note that this only has an effect when using a chunk size
+                # above.
                 if (
                     torch.rand((1,)).item()
                     < self.hparams.dynamic_left_context_prob
diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 8ddc9fa811..e424bd05fd 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -27,14 +27,38 @@
 
 @dataclass
 class ConformerEncoderLayerStreamingContext:
+    """Streaming metadata and state for a `ConformerEncoderLayer`.
+
+    The multi-head attention and Dynamic Chunk Convolution require to save some
+    left context that gets inserted as left padding."""
+
     mha_left_context_size: int
-    mha_left_context: Optional[torch.Tensor]
-    dcconv_left_context: Optional[torch.Tensor]
+    """For this layer, specifies how many frames of inputs should be saved.
+    Usually, the same value is used across all layers, but this can be modified.
+    """
+
+    mha_left_context: Optional[torch.Tensor] = None
+    """Left context to insert at the left of the current chunk as inputs to the
+    multi-head attention. It can be `None` (if we're dealing with the first
+    chunk) or `<= mha_left_context_size` because for the first few chunks, not
+    enough left context may be available to pad.
+    """
+
+    dcconv_left_context: Optional[torch.Tensor] = None
+    """Left context to insert at the left of the convolution according to the
+    Dynamic Chunk Convolution method.
+    
+    Unlike `mha_left_context`, here the amount of frames to keep is fixed and
+    inferred from the kernel size of the convolution module.
+    """
 
 
 @dataclass
 class ConformerEncoderStreamingContext:
+    """Streaming metadata and state for a `ConformerEncoder`."""
+
     layers: List[ConformerEncoderLayerStreamingContext]
+    """Streaming metadata and state for each layer of the encoder."""
 
 
 class ConvolutionModule(nn.Module):
@@ -406,32 +430,25 @@ def forward(
         x = self.norm2(x + 0.5 * self.ffn_module2(x))
         return x, self_attn
 
-    def streaming_forward(
+    def forward_streaming(
         self,
         x,
         context: ConformerEncoderLayerStreamingContext,
         pos_embs: torch.Tensor = None,
     ):
         orig_len = x.shape[-2]
-        # print("pre ffn x:  ", x.shape)
         # ffn module
         x = x + 0.5 * self.ffn_module1(x)
 
-        # print("x:          ", x.shape)
         if context.mha_left_context is not None:
             x = torch.cat((context.mha_left_context, x), dim=1)
 
-        # print("cat(lc, x): ", x.shape)
         if context.mha_left_context_size > 0:
             context.mha_left_context = x[
                 ..., -context.mha_left_context_size :, :
             ]
 
-        # print("pos_embs:   ", pos_embs.shape)
-        # print("new lc:     ", context.mha_left_context.shape)
-        # print()
-
-        # muti-head attention module
+        # multi-head attention module
         skip = x
         x = self.norm1(x)
 
@@ -458,10 +475,16 @@ def streaming_forward(
         return x, self_attn
 
     def make_streaming_context(self, mha_left_context_size: int):
+        """Creates a blank streaming context for this encoding layer.
+
+        Arguments
+        ---------
+        mha_left_context_size : int
+            How many left frames should be saved and used as left context to the
+            current chunk when streaming
+        """
         return ConformerEncoderLayerStreamingContext(
-            mha_left_context_size=mha_left_context_size,
-            mha_left_context=None,
-            dcconv_left_context=None,
+            mha_left_context_size=mha_left_context_size
         )
 
 
@@ -606,7 +629,7 @@ def forward_streaming(
         output = src
         attention_lst = []
         for i, enc_layer in enumerate(self.layers):
-            output, attention = enc_layer.streaming_forward(
+            output, attention = enc_layer.forward_streaming(
                 output, pos_embs=pos_embs, context=context.layers[i]
             )
             attention_lst.append(attention)
@@ -615,6 +638,15 @@ def forward_streaming(
         return output, attention_lst
 
     def make_streaming_context(self, mha_left_context_size: int):
+        """Creates a blank streaming context for the encoder.
+
+        Arguments
+        ---------
+        mha_left_context_size : int
+            How many left frames should be saved and used as left context to the
+            current chunk when streaming. This value is replicated across all
+            layers.
+        """
         return ConformerEncoderStreamingContext(
             layers=[
                 layer.make_streaming_context(
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index b8716772ad..5b99f1c639 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -22,9 +22,16 @@
 
 @dataclass
 class TransformerASRStreamingContext:
+    """Streaming metadata and state for a `TransformerASR` instance."""
+
     chunk_size: int
-    left_context_target_size: int
+    """The size of a chunk expressed in input frames to the transformer."""
+
     encoder_context: Any
+    """Opaque encoder context information. It is constructed by the encoder's
+    `make_streaming_context` method and is passed to the encoder when using
+    `encode_streaming`.
+    """
 
 
 class TransformerASR(TransformerInterface):
@@ -449,14 +456,24 @@ def encode_streaming(self, src, context: TransformerASRStreamingContext):
         return encoder_out
 
     def make_streaming_context(
-        self, chunk_size, left_context_size, encoder_kwargs={}
+        self, chunk_size: int, encoder_kwargs={}
     ):
+        """Creates a blank streaming context for this transformer and its
+        encoder.
+        
+        Arguments
+        ---------
+        chunk_size : int
+            How many frames comprise a chunk.
+        
+        encoder_kwargs : dict
+            Parameters to be forward to the encoder's `make_streaming_context`.
+            Metadata required for the encoder could differ depending on the
+            encoder.
+        """
         return TransformerASRStreamingContext(
             chunk_size=chunk_size,
-            left_context_target_size=left_context_size,
             encoder_context=self.encoder.make_streaming_context(
-                # FIXME: bad naming, not all encoders might use mha etc
-                mha_left_context_size=left_context_size,
                 **encoder_kwargs,
             ),
         )

From a256508dd10ae21b84df7c6e53c48fd966501614 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 23 Aug 2023 15:20:00 +0200
Subject: [PATCH 14/83] More comments and TODOs

---
 .../lobes/models/transformer/Conformer.py     | 33 +++++++++++++++----
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index e424bd05fd..911ef77171 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -130,10 +130,11 @@ def __init__(
             bias=bias,
         )
 
-        # self.batch_norm = nn.BatchNorm1d(input_size)
+        # NOTE: there appears to be a mismatch compared to the Conformer paper:
+        # I believe the first LayerNorm below is supposed to be a BatchNorm.
 
         self.after_conv = nn.Sequential(
-            nn.LayerNorm(input_size),  # should be a BN to match Conformer paper
+            nn.LayerNorm(input_size),
             activation(),
             # pointwise
             nn.Linear(input_size, input_size, bias=bias),
@@ -151,8 +152,8 @@ def _do_conv(self, x, inhibit_padding: bool):
             # let's keep backwards compat by pointing at the weights from the
             # already declared Conv1d.
 
-            # we do not need to edit bottleneck as it is pointwise (i.e. time
-            # step by time step), thus, it doesn't need padding along the
+            # we do not need to edit the bottleneck as it is pointwise (i.e.
+            # time step by time step), thus, it doesn't need padding along the
             # time dimension
             out = F.conv1d(
                 out,
@@ -168,8 +169,6 @@ def _do_conv(self, x, inhibit_padding: bool):
             # chomp
             out = out[..., : -self.padding]
 
-        # out = self.batch_norm(out)
-
         out = out.transpose(1, 2)
         out = self.after_conv(out)
         return out
@@ -241,7 +240,8 @@ def forward(self, x, mask=None, chunk_size=-1):
             ]
 
             # we pack together chunks in a single tensor so that we can feed it
-            # to the convolution directly.
+            # to the convolution directly. this is much more performant than
+            # doing the same with lists.
 
             # -> [batch_size, num_chunks, chunk_size + lc + rpad, in_channels]
             out = torch.stack(out, dim=1)
@@ -440,9 +440,17 @@ def forward_streaming(
         # ffn module
         x = x + 0.5 * self.ffn_module1(x)
 
+        # TODO: make the approach for MHA left context more efficient.
+        # currently, this saves the inputs to the MHA.
+        # the naive approach is suboptimal in a few ways, namely that the
+        # outputs for this left padding is being re-computed even though we
+        # discard them immediately after.
+
+        # left pad `x` with our MHA left context
         if context.mha_left_context is not None:
             x = torch.cat((context.mha_left_context, x), dim=1)
 
+        # compute new MHA left context for the next call to our function
         if context.mha_left_context_size > 0:
             context.mha_left_context = x[
                 ..., -context.mha_left_context_size :, :
@@ -456,11 +464,21 @@ def forward_streaming(
             x, x, x, attn_mask=None, key_padding_mask=None, pos_embs=pos_embs,
         )
         x = x + skip
+
+        # truncate outputs corresponding to the MHA left context (we only care
+        # about our chunk's outputs); see above to-do
         x = x[..., -orig_len:, :]
 
+        # TODO: this is slightly suboptimal as this will add left padding inside
+        # the convolution code that we do not need. it would be better to
+        # manually add the right-padding ourselves and disable padding inside
+        # the convolution module for this usecase, but it would need some
+        # refactoring.
+
         if context.dcconv_left_context is not None:
             x = torch.cat((context.dcconv_left_context, x), dim=1)
 
+        # compute new DCConv left context for the next call to our function
         context.dcconv_left_context = x[
             ..., -self.convolution_module.padding :, :
         ]
@@ -468,6 +486,7 @@ def forward_streaming(
         # convolution module
         x = x + self.convolution_module(x)
 
+        # truncate outputs corresponding to the DCConv left context
         x = x[..., -orig_len:, :]
 
         # ffn module

From 0e69129c63740854bc74852e0acc07fd8b2395a1 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 23 Aug 2023 15:24:31 +0200
Subject: [PATCH 15/83] encode_streaming docstring

---
 .../lobes/models/transformer/TransformerASR.py    | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 5b99f1c639..57103bddbf 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -411,8 +411,21 @@ def encode(
     def encode_streaming(self, src, context: TransformerASRStreamingContext):
         """
         Streaming encoder forward pass
+
+        Arguments
+        ---------
+        src : torch.Tensor
+            The sequence (chunk) to the encoder.
+
+        context : TransformerASRStreamingContext
+            Mutable reference to the streaming context. This holds the state
+            needed to persist across chunk inferences and can be built using
+            `make_streaming_context`. This will get mutated by this function.
+
+        Returns
+        -------
+        Encoder output for this chunk.
         """
-        # TODO: docstring
 
         if src.dim() == 4:
             bz, t, ch1, ch2 = src.shape

From d05a771e845944959af0a70bd0687934194e1bdd Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 28 Aug 2023 16:05:16 +0200
Subject: [PATCH 16/83] Dirty TransducerBeamSearcher change for streaming GS

---
 speechbrain/decoders/transducer.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/speechbrain/decoders/transducer.py b/speechbrain/decoders/transducer.py
index 9cc9db400c..87be90af46 100644
--- a/speechbrain/decoders/transducer.py
+++ b/speechbrain/decoders/transducer.py
@@ -135,7 +135,7 @@ def forward(self, tn_output):
         hyps = self.searcher(tn_output)
         return hyps
 
-    def transducer_greedy_decode(self, tn_output):
+    def transducer_greedy_decode(self, tn_output, start_state=None, return_hidden=False):
         """Transducer greedy decoder is a greedy decoder over batch which apply Transducer rules:
             1- for each time step in the Transcription Network (TN) output:
                 -> Update the ith utterance only if
@@ -160,7 +160,6 @@ def transducer_greedy_decode(self, tn_output):
             "logp_scores": [0.0 for _ in range(tn_output.size(0))],
         }
         # prepare BOS = Blank for the Prediction Network (PN)
-        hidden = None
         input_PN = (
             torch.ones(
                 (tn_output.size(0), 1),
@@ -169,8 +168,13 @@ def transducer_greedy_decode(self, tn_output):
             )
             * self.blank_id
         )
-        # First forward-pass on PN
-        out_PN, hidden = self._forward_PN(input_PN, self.decode_network_lst)
+
+        if start_state is None:
+            # First forward-pass on PN
+            out_PN, hidden = self._forward_PN(input_PN, self.decode_network_lst)
+        else:
+            out_PN, hidden = start_state
+
         # 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]
@@ -210,13 +214,19 @@ def transducer_greedy_decode(self, tn_output):
                     have_update_hyp, selected_hidden, hidden
                 )
 
-        return (
+        ret = (
             hyp["prediction"],
             torch.Tensor(hyp["logp_scores"]).exp().mean(),
             None,
             None,
         )
 
+        if return_hidden:
+            ret += ((out_PN, hidden,),)
+
+        return ret
+
+
     def transducer_beam_search_decode(self, tn_output):
         """Transducer beam search decoder is a beam search decoder over batch which apply Transducer rules:
             1- for each utterance:

From afb96dbfeaa92a83c0f87eee2e6341abdaaee99a Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 13:50:14 +0200
Subject: [PATCH 17/83] Fix precommit

---
 speechbrain/decoders/transducer.py                     | 5 +++--
 speechbrain/lobes/models/transformer/Conformer.py      | 2 +-
 speechbrain/lobes/models/transformer/TransformerASR.py | 8 +++-----
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/speechbrain/decoders/transducer.py b/speechbrain/decoders/transducer.py
index 87be90af46..c5b585cf95 100644
--- a/speechbrain/decoders/transducer.py
+++ b/speechbrain/decoders/transducer.py
@@ -135,7 +135,9 @@ def forward(self, tn_output):
         hyps = self.searcher(tn_output)
         return hyps
 
-    def transducer_greedy_decode(self, tn_output, start_state=None, return_hidden=False):
+    def transducer_greedy_decode(
+        self, tn_output, start_state=None, return_hidden=False
+    ):
         """Transducer greedy decoder is a greedy decoder over batch which apply Transducer rules:
             1- for each time step in the Transcription Network (TN) output:
                 -> Update the ith utterance only if
@@ -226,7 +228,6 @@ def transducer_greedy_decode(self, tn_output, start_state=None, return_hidden=Fa
 
         return ret
 
-
     def transducer_beam_search_decode(self, tn_output):
         """Transducer beam search decoder is a beam search decoder over batch which apply Transducer rules:
             1- for each utterance:
diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 911ef77171..98c9c11310 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -47,7 +47,7 @@ class ConformerEncoderLayerStreamingContext:
     dcconv_left_context: Optional[torch.Tensor] = None
     """Left context to insert at the left of the convolution according to the
     Dynamic Chunk Convolution method.
-    
+
     Unlike `mha_left_context`, here the amount of frames to keep is fixed and
     inferred from the kernel size of the convolution module.
     """
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 57103bddbf..79e0456162 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -468,17 +468,15 @@ def encode_streaming(self, src, context: TransformerASRStreamingContext):
         )
         return encoder_out
 
-    def make_streaming_context(
-        self, chunk_size: int, encoder_kwargs={}
-    ):
+    def make_streaming_context(self, chunk_size: int, encoder_kwargs={}):
         """Creates a blank streaming context for this transformer and its
         encoder.
-        
+
         Arguments
         ---------
         chunk_size : int
             How many frames comprise a chunk.
-        
+
         encoder_kwargs : dict
             Parameters to be forward to the encoder's `make_streaming_context`.
             Metadata required for the encoder could differ depending on the

From 6585ae857ed46e0cca2c39885b34fd16c2841b8b Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 14:14:54 +0200
Subject: [PATCH 18/83] Fix encoders that do not support chunk_size

---
 speechbrain/lobes/models/transformer/Branchformer.py   | 3 +++
 speechbrain/lobes/models/transformer/Transformer.py    | 2 ++
 speechbrain/lobes/models/transformer/TransformerASR.py | 6 ++----
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Branchformer.py b/speechbrain/lobes/models/transformer/Branchformer.py
index 3f75701a34..ad03e6a0c2 100644
--- a/speechbrain/lobes/models/transformer/Branchformer.py
+++ b/speechbrain/lobes/models/transformer/Branchformer.py
@@ -320,6 +320,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
+        chunk_size: int = -1,
     ):
         """
         Arguments
@@ -336,6 +337,8 @@ def forward(
             where S is the sequence length, and E is the embedding dimension.
         """
 
+        assert chunk_size == -1, "Encoder does not support streaming"
+
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:
                 raise ValueError(
diff --git a/speechbrain/lobes/models/transformer/Transformer.py b/speechbrain/lobes/models/transformer/Transformer.py
index 012d87ec44..add6562e7c 100644
--- a/speechbrain/lobes/models/transformer/Transformer.py
+++ b/speechbrain/lobes/models/transformer/Transformer.py
@@ -526,6 +526,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
+        chunk_size: int = -1,
     ):
         """
         Arguments
@@ -537,6 +538,7 @@ def forward(
         src_key_padding_mask : tensor
             The mask for the src keys per batch (optional).
         """
+        assert chunk_size == -1, "Encoder does not support streaming"
         output = src
         if self.layerdrop_prob > 0.0:
             keep_probs = self.rng.random(len(self.layers))
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 79e0456162..dc37826bcb 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -188,7 +188,6 @@ def forward(self, src, tgt, wav_len=None, pad_idx=0):
         pad_idx : int, optional
             The index for <pad> token (default=0).
         """
-        # FIXME: should have chunk_size & left_chunk_context?
 
         # reshpae the src vector to [Batch, Time, Fea] is a 4d vector is given
         if src.ndim == 4:
@@ -525,14 +524,13 @@ def __init__(self, transformer, *args, **kwargs):
         self.transformer = transformer
 
     def forward(
-        self, x, wav_lens=None, pad_idx=0, chunk_size=-1, left_context_chunks=-1
+        self, x, wav_lens=None, pad_idx=0, **kwargs
     ):
         """ Processes the input tensor x and returns an output tensor."""
         x = self.transformer.encode(
             x,
             wav_lens,
             pad_idx,
-            chunk_size=chunk_size,
-            left_context_chunks=left_context_chunks,
+            **kwargs,
         )
         return x

From 98d0ddfde36754dadd624fa33e5ef68f19957568 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 14:16:52 +0200
Subject: [PATCH 19/83] Pre-commit again

---
 .../lobes/models/transformer/TransformerASR.py        | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index dc37826bcb..bd2bc42161 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -523,14 +523,7 @@ def __init__(self, transformer, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.transformer = transformer
 
-    def forward(
-        self, x, wav_lens=None, pad_idx=0, **kwargs
-    ):
+    def forward(self, x, wav_lens=None, pad_idx=0, **kwargs):
         """ Processes the input tensor x and returns an output tensor."""
-        x = self.transformer.encode(
-            x,
-            wav_lens,
-            pad_idx,
-            **kwargs,
-        )
+        x = self.transformer.encode(x, wav_lens, pad_idx, **kwargs,)
         return x

From c23435f9b424c7383e03cf70237f6097e12f3f68 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 14:29:20 +0200
Subject: [PATCH 20/83] Make chunk_size type consistent

---
 speechbrain/lobes/models/transformer/Conformer.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 98c9c11310..fb5a7b0a86 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -386,7 +386,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: torch.Tensor = None,
-        chunk_size: Optional[int] = None,
+        chunk_size: int = -1,
     ):
         """
         Arguments
@@ -399,7 +399,7 @@ def forward(
             The mask for the src keys per batch.
         pos_embs: torch.Tensor, torch.nn.Module, optional
             Module or tensor containing the input sequence positional embeddings
-        chunk_size: int, optional
+        chunk_size: int
             Whether to preform convolution chunking to hide future context,
             useful for chunked conformers in a dynamic chunk training setting
         """
@@ -593,7 +593,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
-        chunk_size: Optional[int] = None,
+        chunk_size: int = -1,
     ):
         """
         Arguments
@@ -608,9 +608,10 @@ def forward(
             Module or tensor containing the input sequence positional embeddings
             If custom pos_embs are given it needs to have the shape (1, 2*S-1, E)
             where S is the sequence length, and E is the embedding dimension.
-        chunk_size: int, optional
+        chunk_size: int
             Whether to preform convolution chunking to hide future context,
             useful for chunked conformers in a dynamic chunk training setting
+            `-1` ignores chunking
         """
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:

From dd264a6d00ab5b4c00d9617280b1313623c05052 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 14:38:35 +0200
Subject: [PATCH 21/83] Fix formatting of doctest in split_wav_lens

---
 speechbrain/utils/streaming.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/speechbrain/utils/streaming.py b/speechbrain/utils/streaming.py
index 15911db451..336f6e197a 100644
--- a/speechbrain/utils/streaming.py
+++ b/speechbrain/utils/streaming.py
@@ -96,9 +96,7 @@ def split_wav_lens(
     >>> wav_lens = torch.tensor([1.0, 0.65, 0.85])
     >>> chunked_wav_lens = split_wav_lens([c.size(1) for c in chunks], wav_lens)
     >>> chunked_wav_lens
-    [tensor([1., 1., 1.]),
-     tensor([1.0000, 0.6250, 1.0000]),
-     tensor([1.0000, 0.0000, 0.2500])]
+    [tensor([1., 1., 1.]), tensor([1.0000, 0.6250, 1.0000]), tensor([1.0000, 0.0000, 0.2500])]
     >>> # wav 1 covers 62.5% (5/8) of the second chunk's frames
     """
 

From 107688eda601914b26013eea295d394dc3899a4d Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 14:39:15 +0200
Subject: [PATCH 22/83] Remove outdated TODO

---
 .../ASR/transducer/hparams/conformer_transducer.yaml            | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index a4957844a3..c386e014ea 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -87,7 +87,7 @@ dynamic_chunk_prob: 0.6
 dynamic_chunk_min: 8
 dynamic_chunk_max: 32
 
-dynamic_left_context_prob: 0.75 # TODO: rename all to prob, makes more sense
+dynamic_left_context_prob: 0.75
 dynamic_left_context_min: 16
 dynamic_left_context_max: 64
 

From 8b88dc9866305ff6acc7365fd21633dbc6cbdce9 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 14:50:24 +0200
Subject: [PATCH 23/83] Add hasattr streaming to retain model backcompat

---
 recipes/LibriSpeech/ASR/transducer/train.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 2b5eaae8e3..efc903c8c5 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -73,7 +73,10 @@ def compute_forward(self, batch, stage):
         # Default to infinite context visibility
         transformer_chunk_size = -1
         left_context_chunks = -1
-        if self.hparams.streaming:
+
+        # Old models may not have the streaming hparam, we don't break them in
+        # any other way so just check for its presence
+        if hasattr(self.hparams, "streaming") and self.hparams.streaming:
             # TODO: while this is fairly small logic, it may make sense to
             # extract it to its own class orchestrating dynamic chunk training,
             # partly because explantions are beneficial

From c4c730da8c48e26bc6ba8e3a4071b427f0bbed7a Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 15:17:55 +0200
Subject: [PATCH 24/83] Cleanup doc and naming for transducer_greedy_decode

---
 speechbrain/decoders/transducer.py | 35 ++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/speechbrain/decoders/transducer.py b/speechbrain/decoders/transducer.py
index c5b585cf95..b96eae9ff2 100644
--- a/speechbrain/decoders/transducer.py
+++ b/speechbrain/decoders/transducer.py
@@ -136,7 +136,7 @@ def forward(self, tn_output):
         return hyps
 
     def transducer_greedy_decode(
-        self, tn_output, start_state=None, return_hidden=False
+        self, tn_output, hidden_state=None, return_hidden=False
     ):
         """Transducer greedy decoder is a greedy decoder over batch which apply Transducer rules:
             1- for each time step in the Transcription Network (TN) output:
@@ -151,11 +151,37 @@ def transducer_greedy_decode(
             Output from transcription network with shape
             [batch, time_len, hiddens].
 
+        hidden_state : (torch.Tensor, torch.Tensor)
+            Hidden state to initially feed the decode network with. This is
+            useful in conjunction with `return_hidden` to be able to perform
+            beam search in a streaming context, so that you can reuse the last
+            hidden state as an initial state across calls.
+
+        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`.
+
         Returns
         -------
-        torch.tensor
+        Tuple of 4 or 5 elements (if `return_hidden`).
+
+        First element: List[List[int]]
+            List of decoded tokens
+
+        Second element: torch.Tensor
             Outputs a logits tensor [B,T,1,Output_Dim]; padding
             has not been removed.
+
+        Third element: None
+            nbest; irrelevant for greedy decode
+
+        Fourth element: None
+            nbest scores; irrelevant for greedy decode
+
+        Fifth element: Present if `return_hidden`, (torch.Tensor, torch.Tensor)
+            Tuple representing the hidden state required to call
+            `transducer_greedy_decode` where you left off in a streaming
+            context.
         """
         hyp = {
             "prediction": [[] for _ in range(tn_output.size(0))],
@@ -171,11 +197,11 @@ def transducer_greedy_decode(
             * self.blank_id
         )
 
-        if start_state is None:
+        if hidden_state is None:
             # First forward-pass on PN
             out_PN, hidden = self._forward_PN(input_PN, self.decode_network_lst)
         else:
-            out_PN, hidden = start_state
+            out_PN, hidden = hidden_state
 
         # For each time step
         for t_step in range(tn_output.size(1)):
@@ -224,6 +250,7 @@ def transducer_greedy_decode(
         )
 
         if return_hidden:
+            # append the `(out_PN, hidden)` tuple to ret
             ret += ((out_PN, hidden,),)
 
         return ret

From a02ed5f8ce30fc93be19e024ea614e5708693062 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 15:29:56 +0200
Subject: [PATCH 25/83] Cite paper for chunked attention

---
 speechbrain/lobes/models/transformer/Conformer.py      | 1 -
 speechbrain/lobes/models/transformer/TransformerASR.py | 5 +++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index fb5a7b0a86..00edb39d97 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -403,7 +403,6 @@ def forward(
             Whether to preform convolution chunking to hide future context,
             useful for chunked conformers in a dynamic chunk training setting
         """
-        # TODO: cite paper for chunk size
         # TODO: document left frames
 
         conv_mask: Optional[torch.Tensor] = None
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index bd2bc42161..ce4b2a5974 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -292,6 +292,11 @@ def make_masks(
                 # end range is exclusive, so there is no off-by-one here
                 src_mask[i, :frame_remaining_context] = True
 
+        # The following is not really the sole source used to implement this,
+        # but it introduces the concept.
+        # ref: Unified Streaming and Non-streaming Two-pass End-to-end Model
+        # for Speech Recognition
+        # https://arxiv.org/pdf/2012.05481.pdf
         if chunk_size >= 0:
             for i in range(src.shape[1]):
                 # if we have a chunk size of 8 then:

From be92a12c27233dfe6b7687e9ed22797e9c809e89 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 15:33:41 +0200
Subject: [PATCH 26/83] Remove lost comment

---
 speechbrain/lobes/models/transformer/Conformer.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 00edb39d97..f890b221f1 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -403,8 +403,6 @@ def forward(
             Whether to preform convolution chunking to hide future context,
             useful for chunked conformers in a dynamic chunk training setting
         """
-        # TODO: document left frames
-
         conv_mask: Optional[torch.Tensor] = None
         if src_key_padding_mask is not None:
             conv_mask = src_key_padding_mask.unsqueeze(-1)

From 382b97b60bfda774945b7fe3feb2d9c69924e185 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 15:49:20 +0200
Subject: [PATCH 27/83] Update comment in self-attention

---
 speechbrain/nnet/attention.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/speechbrain/nnet/attention.py b/speechbrain/nnet/attention.py
index 74d1305a6f..0f983634be 100644
--- a/speechbrain/nnet/attention.py
+++ b/speechbrain/nnet/attention.py
@@ -591,11 +591,12 @@ def forward(
             query + self.pos_bias_v.view(1, 1, self.num_heads, self.head_dim)
         ).transpose(1, 2)
 
-        # TODO: cite https://asherliu.github.io/docs/sc21a.pdf
-        # for the scaling prior to the matrix multiplication
-        # should read more of the paper though
-        # TODO: check if this causes any difference beyond precision
-        # (it does not seem like it does)
+        # Moved the `* self.scale` mul from after the `attn_score` sum to prior
+        # to the matmul in order to lower overflow risks on fp16.
+        # This change is inspired by the following paper, but no other changes
+        # were ported from there so far.
+        # ref: E.T.: Re-Thinking Self-Attention for Transformer Models on GPUs
+        # https://asherliu.github.io/docs/sc21a.pdf
 
         # (batch, head, qlen, klen)
         matrix_ac = torch.matmul(

From 12f89bf14090e62aa07346199963441b27635d8f Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 15:53:32 +0200
Subject: [PATCH 28/83] Don't apply masked fill fix in the non-bool mask case

---
 speechbrain/nnet/attention.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/speechbrain/nnet/attention.py b/speechbrain/nnet/attention.py
index 0f983634be..f6351dd690 100644
--- a/speechbrain/nnet/attention.py
+++ b/speechbrain/nnet/attention.py
@@ -643,7 +643,9 @@ def forward(
             if attn_mask.dtype == torch.bool:
                 attn_score = attn_score.masked_fill(attn_mask, 0.0)
             else:
-                assert False, "oopsie need to reimplement that"
+                # NOTE: the above fix is not implemented for this case as
+                # summing the mask with NaN would still result in NaN
+                pass
 
         if key_padding_mask is not None:
             attn_score = attn_score.masked_fill(

From ee444a0b1457dfd21d855c5c7d18b1026fa6fc69 Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Tue, 29 Aug 2023 16:13:29 +0200
Subject: [PATCH 29/83] Added TODO README update

---
 recipes/LibriSpeech/ASR/transducer/README.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/README.md b/recipes/LibriSpeech/ASR/transducer/README.md
index f3e6d3d1ea..402824edcf 100644
--- a/recipes/LibriSpeech/ASR/transducer/README.md
+++ b/recipes/LibriSpeech/ASR/transducer/README.md
@@ -26,7 +26,8 @@ Dev. clean is evaluated with Greedy Decoding while the test sets are using Greed
 
 | Release | Hyperparams file | Dev-clean Greedy | Test-clean Greedy | Test-other Greedy | Test-clean BS+RNNLM | Test-other BS+RNNLM | Model link | GPUs |
 |:-------------:|:---------------------------:| :------:| :-----------:| :------------------:| :------------------:| :------------------:| :--------:| :-----------:|
-| 2023-07-19 | conformer_transducer.yaml | 2.62 | 2.84 | 6.98 | 2.62 | 6.31 | https://drive.google.com/drive/folders/1QtQz1Bkd_QPYnf3CyxhJ57ovbSZC2EhN?usp=sharing | 3x3090 24GB |
+| 2023-07-19 | conformer_transducer.yaml `streaming: False` | 2.62 | 2.84 | 6.98 | 2.62 | 6.31 | https://drive.google.com/drive/folders/1QtQz1Bkd_QPYnf3CyxhJ57ovbSZC2EhN?usp=sharing | 3x3090 24GB |
+| 2023-08-29 | conformer_transducer.yaml `streaming: True`  | TODO | TODO | TODO | Untested | Untested | TODO | 1x A100 40GB(?) |
 
 # **About SpeechBrain**
 - Website: https://speechbrain.github.io/

From 1013c717a8c77d1acdb8d982ce54cfb684cdb2da Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Wed, 30 Aug 2023 10:57:47 +0200
Subject: [PATCH 30/83] Revert change to custom_tgt_module; patching model
 instead

---
 speechbrain/lobes/models/transformer/TransformerASR.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index ce4b2a5974..6c1c5f60ae 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -167,10 +167,9 @@ def __init__(
             ),
             torch.nn.Dropout(dropout),
         )
-        if num_decoder_layers > 0:
-            self.custom_tgt_module = ModuleList(
-                NormalizedEmbedding(d_model, tgt_vocab)
-            )
+        self.custom_tgt_module = ModuleList(
+            NormalizedEmbedding(d_model, tgt_vocab)
+        )
 
         # reset parameters using xavier_normal_
         self._init_params()

From 10ff2159eda9afe6f3d5ff3171cacecf137de0dc Mon Sep 17 00:00:00 2001
From: Sylvain de Langen <sdelang@sdelang.fr>
Date: Thu, 31 Aug 2023 13:57:33 +0200
Subject: [PATCH 31/83] Remove added entry in README

---
 recipes/LibriSpeech/ASR/transducer/README.md | 1 -
 1 file changed, 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/README.md b/recipes/LibriSpeech/ASR/transducer/README.md
index 402824edcf..6fad3a0059 100644
--- a/recipes/LibriSpeech/ASR/transducer/README.md
+++ b/recipes/LibriSpeech/ASR/transducer/README.md
@@ -27,7 +27,6 @@ Dev. clean is evaluated with Greedy Decoding while the test sets are using Greed
 | Release | Hyperparams file | Dev-clean Greedy | Test-clean Greedy | Test-other Greedy | Test-clean BS+RNNLM | Test-other BS+RNNLM | Model link | GPUs |
 |:-------------:|:---------------------------:| :------:| :-----------:| :------------------:| :------------------:| :------------------:| :--------:| :-----------:|
 | 2023-07-19 | conformer_transducer.yaml `streaming: False` | 2.62 | 2.84 | 6.98 | 2.62 | 6.31 | https://drive.google.com/drive/folders/1QtQz1Bkd_QPYnf3CyxhJ57ovbSZC2EhN?usp=sharing | 3x3090 24GB |
-| 2023-08-29 | conformer_transducer.yaml `streaming: True`  | TODO | TODO | TODO | Untested | Untested | TODO | 1x A100 40GB(?) |
 
 # **About SpeechBrain**
 - Website: https://speechbrain.github.io/

From b16754f7f97685a69e4126477ca8dc5e940a5518 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 20 Nov 2023 13:26:17 +0100
Subject: [PATCH 32/83] Fix streaming conformer conv mismatch

---
 .../lobes/models/transformer/Conformer.py     | 22 ++++++++++++-------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index f890b221f1..d58d28b78b 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -142,12 +142,8 @@ def __init__(
         )
 
     def _do_conv(self, x, inhibit_padding: bool):
-        out = self.layer_norm(x)
-        out = out.transpose(1, 2)
-        out = self.bottleneck(out)
-
         if not inhibit_padding:
-            out = self.conv(out)
+            out = self.conv(x)
         else:
             # let's keep backwards compat by pointing at the weights from the
             # already declared Conv1d.
@@ -156,13 +152,13 @@ def _do_conv(self, x, inhibit_padding: bool):
             # time step by time step), thus, it doesn't need padding along the
             # time dimension
             out = F.conv1d(
-                out,
+                x,
                 weight=self.conv.weight,
                 bias=self.conv.bias,
                 stride=self.conv.stride,
                 padding=0,
                 dilation=self.conv.dilation,
-                groups=out.shape[-2],
+                groups=x.shape[-2],
             )
 
         if self.causal:
@@ -219,6 +215,11 @@ def forward(self, x, mask=None, chunk_size=-1):
                 for i in range(chunk_count)
             ]
 
+            out = [self.layer_norm(chk) for chk in out]
+            out = [chk.transpose(1, 2) for chk in out]
+            out = [self.bottleneck(chk) for chk in out]
+            out = [chk.transpose(1, 2) for chk in out]
+
             # TODO: experiment around reflect padding, which is difficult
             # because small chunks have too little time steps to reflect from
             out = [
@@ -249,6 +250,8 @@ def forward(self, x, mask=None, chunk_size=-1):
             # -> [batch_size * num_chunks, chunk_size + lc + rpad, in_channels]
             out = torch.flatten(out, end_dim=1)
 
+            out = out.transpose(1, 2)
+
             # -> [batch_size * num_chunks, chunk_size, out_channels]
             out = self._do_conv(out, inhibit_padding=True)
 
@@ -262,7 +265,10 @@ def forward(self, x, mask=None, chunk_size=-1):
             if final_right_padding > 0:
                 out = out[:, :-final_right_padding, :]
         else:
-            out = self._do_conv(x, inhibit_padding=False)
+            out = self.layer_norm(x)
+            out = out.transpose(1, 2)
+            out = self.bottleneck(out)
+            out = self._do_conv(out, inhibit_padding=False)
 
         if mask is not None:
             out.masked_fill_(mask, 0.0)

From e5785d89421838e3caab7d6cda1761f875acd27e Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Tue, 21 Nov 2023 10:15:01 +0100
Subject: [PATCH 33/83] More conformer conv adjustments

---
 .../lobes/models/transformer/Conformer.py     | 25 +++++++++++++------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index d58d28b78b..646c27c58f 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -158,15 +158,9 @@ def _do_conv(self, x, inhibit_padding: bool):
                 stride=self.conv.stride,
                 padding=0,
                 dilation=self.conv.dilation,
-                groups=x.shape[-2],
+                groups=self.conv.groups,
             )
 
-        if self.causal:
-            # chomp
-            out = out[..., : -self.padding]
-
-        out = out.transpose(1, 2)
-        out = self.after_conv(out)
         return out
 
     def forward(self, x, mask=None, chunk_size=-1):
@@ -236,6 +230,7 @@ def forward(self, x, mask=None, chunk_size=-1):
                         self.padding
                         + (final_right_padding if i == len(out) - 1 else 0),
                     ),
+                    value=0
                 )
                 for i in range(len(out))
             ]
@@ -250,11 +245,18 @@ def forward(self, x, mask=None, chunk_size=-1):
             # -> [batch_size * num_chunks, chunk_size + lc + rpad, in_channels]
             out = torch.flatten(out, end_dim=1)
 
+            # for the convolution:
+            # -> [batch_size * num_chunks, in_channels, chunk_size + lc + rpad]
             out = out.transpose(1, 2)
 
-            # -> [batch_size * num_chunks, chunk_size, out_channels]
+            # -> [batch_size * num_chunks, out_channels, chunk_size + rpad]
             out = self._do_conv(out, inhibit_padding=True)
 
+            # -> [batch_size * num_chunks, chunk_size + rpad, out_channels]
+            out = out.transpose(1, 2)
+
+            out = self.after_conv(out)
+
             # -> [batch_size, num_chunks, chunk_size, out_channels]
             out = torch.unflatten(out, dim=0, sizes=(batch_size, -1))
 
@@ -270,6 +272,13 @@ def forward(self, x, mask=None, chunk_size=-1):
             out = self.bottleneck(out)
             out = self._do_conv(out, inhibit_padding=False)
 
+            out = out.transpose(1, 2)
+            out = self.after_conv(out)
+
+            if self.causal:
+                # chomp
+                out = out[..., : -self.padding]
+
         if mask is not None:
             out.masked_fill_(mask, 0.0)
 

From 96331564a6afc713f7432d37b6e74392f7adaadb Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Tue, 21 Nov 2023 14:26:48 +0100
Subject: [PATCH 34/83] Adjust context size

---
 .../ASR/transducer/hparams/conformer_transducer.yaml         | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index c386e014ea..a986184b47 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -87,9 +87,10 @@ dynamic_chunk_prob: 0.6
 dynamic_chunk_min: 8
 dynamic_chunk_max: 32
 
+# Left context gets expressed as a number of chunks
 dynamic_left_context_prob: 0.75
-dynamic_left_context_min: 16
-dynamic_left_context_max: 64
+dynamic_left_context_min: 2
+dynamic_left_context_max: 32
 
 # Dataloader options
 train_dataloader_opts:

From c1fbb8f0fcb56ec0bb640b18c79a2ca0d14ac421 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Thu, 23 Nov 2023 15:17:59 +0100
Subject: [PATCH 35/83] Remove outdated comment

---
 recipes/LibriSpeech/ASR/transducer/train.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index efc903c8c5..585a0f622a 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -58,7 +58,6 @@ def compute_forward(self, batch, stage):
                     tokens_with_bos
                 )
 
-        # Forward pass
         feats = self.hparams.compute_features(wavs)
 
         # Add feature augmentation if specified.

From 74467069a91f667d989e037376365514de527072 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Tue, 28 Nov 2023 12:54:15 +0100
Subject: [PATCH 36/83] Fixed causal conformer decoder

---
 speechbrain/lobes/models/transformer/Conformer.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 646c27c58f..c0743144ae 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -272,13 +272,13 @@ def forward(self, x, mask=None, chunk_size=-1):
             out = self.bottleneck(out)
             out = self._do_conv(out, inhibit_padding=False)
 
-            out = out.transpose(1, 2)
-            out = self.after_conv(out)
-
             if self.causal:
                 # chomp
                 out = out[..., : -self.padding]
 
+            out = out.transpose(1, 2)
+            out = self.after_conv(out)
+
         if mask is not None:
             out.masked_fill_(mask, 0.0)
 

From 1f91e85d254978c32b29f0c351ce8fd700d44efe Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Tue, 28 Nov 2023 13:08:43 +0100
Subject: [PATCH 37/83] Fix linting

---
 speechbrain/lobes/models/transformer/Conformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index c0743144ae..0ce64f2098 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -230,7 +230,7 @@ def forward(self, x, mask=None, chunk_size=-1):
                         self.padding
                         + (final_right_padding if i == len(out) - 1 else 0),
                     ),
-                    value=0
+                    value=0,
                 )
                 for i in range(len(out))
             ]

From d96a92e3bb54bc6294216d67d25c9d26e475f8df Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 4 Dec 2023 14:12:27 +0100
Subject: [PATCH 38/83] Gate `custom_tgt_module` creation behind the presence
 of decoder layers

---
 speechbrain/lobes/models/transformer/TransformerASR.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 6c1c5f60ae..6dfdc31eaa 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -167,9 +167,11 @@ def __init__(
             ),
             torch.nn.Dropout(dropout),
         )
-        self.custom_tgt_module = ModuleList(
-            NormalizedEmbedding(d_model, tgt_vocab)
-        )
+
+        if num_decoder_layers > 0:
+            self.custom_tgt_module = ModuleList(
+                NormalizedEmbedding(d_model, tgt_vocab)
+            )
 
         # reset parameters using xavier_normal_
         self._init_params()

From ddb6d5b1076f8c5d57148f5017b57b5bc6b4574b Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 4 Dec 2023 14:17:52 +0100
Subject: [PATCH 39/83] Re-enable checkpoint averaging

---
 .../ASR/transducer/hparams/conformer_transducer.yaml            | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index a986184b47..f4bc6fb895 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -63,7 +63,7 @@ precision: fp32 # bf16, fp16 or fp32
 batch_size: 8
 grad_accumulation_factor: 4
 sorting: random
-avg_checkpoints: 1 # Number of checkpoints to average for evaluation
+avg_checkpoints: 5 # Number of checkpoints to average for evaluation
 
 # Feature parameters
 sample_rate: 16000

From 6ce59c375b2a9d7bd368ac915d4b81fc2203f2db Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 4 Dec 2023 14:37:52 +0100
Subject: [PATCH 40/83] Change averaged ckpt count to 10

---
 .../ASR/transducer/hparams/conformer_transducer.yaml            | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index f4bc6fb895..fd426c853a 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -63,7 +63,7 @@ precision: fp32 # bf16, fp16 or fp32
 batch_size: 8
 grad_accumulation_factor: 4
 sorting: random
-avg_checkpoints: 5 # Number of checkpoints to average for evaluation
+avg_checkpoints: 10 # Number of checkpoints to average for evaluation
 
 # Feature parameters
 sample_rate: 16000

From 4f52a6f2deb44c0188cc8831a28133233c800456 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Thu, 14 Dec 2023 15:02:53 +0100
Subject: [PATCH 41/83] Add new model results to README

---
 recipes/LibriSpeech/ASR/transducer/README.md | 39 ++++++++++++++++++--
 1 file changed, 35 insertions(+), 4 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/README.md b/recipes/LibriSpeech/ASR/transducer/README.md
index 6fad3a0059..5a28e6965f 100644
--- a/recipes/LibriSpeech/ASR/transducer/README.md
+++ b/recipes/LibriSpeech/ASR/transducer/README.md
@@ -22,11 +22,42 @@ python train.py hparams/conformer_transducer.yaml
 
 # Librispeech Results
 
-Dev. clean is evaluated with Greedy Decoding while the test sets are using Greedy Decoding OR a RNNLM + Beam Search.
+Dev. clean is evaluated with Greedy Decoding while the test sets are using Greedy Decoding OR a RNNLM + Beam Search.  
+Evaluation is performed in fp32.
 
-| Release | Hyperparams file | Dev-clean Greedy | Test-clean Greedy | Test-other Greedy | Test-clean BS+RNNLM | Test-other BS+RNNLM | Model link | GPUs |
-|:-------------:|:---------------------------:| :------:| :-----------:| :------------------:| :------------------:| :------------------:| :--------:| :-----------:|
-| 2023-07-19 | conformer_transducer.yaml `streaming: False` | 2.62 | 2.84 | 6.98 | 2.62 | 6.31 | https://drive.google.com/drive/folders/1QtQz1Bkd_QPYnf3CyxhJ57ovbSZC2EhN?usp=sharing | 3x3090 24GB |
+| Release | Hyperparams file | Train precision | Dev-clean Greedy | Test-clean Greedy | Test-other Greedy | Test-clean BS+RNNLM | Test-other BS+RNNLM | Model link | GPUs |
+|:-------------:|:---------------------------:|:-:| :------:| :-----------:| :------------------:| :------------------:| :------------------:| :--------:| :-----------:|
+| 2023-12-12 | conformer_transducer.yaml `streaming: True` | bf16 | 2.56% | 2.72% | 6.47% | TBD | TBD | https://drive.google.com/drive/folders/1QtQz1Bkd_QPYnf3CyxhJ57ovbSZC2EhN?usp=sharing | [4x A100SXM4 40GB](https://docs.alliancecan.ca/wiki/Narval/en) |
+
+## Streaming model
+
+### WER vs chunk size & left context
+
+**Note:** High-level streaming inference code is not currently available.
+
+The following matrix presents the Word Error Rate (WER%) achieved on LibriSpeech
+`test-clean` with various chunk sizes (in ms) and left context sizes (in # of
+chunks).
+
+The relative difference is not trivial to interpret, because we are not testing
+against a continuous stream of speech, but rather against utterances of various
+lengths. This tends to bias results in favor of larger chunk sizes.
+
+The chunk size might not accurately represent expected latency due to slight
+padding differences in streaming contexts.
+
+The left chunk size is not representative of the receptive field of the model.
+Because the model caches the streaming context at different layers, the model
+may end up forming indirect dependencies to audio many seconds ago.
+
+|       | full | cs=32 (1280ms) | 24 (960ms) | 16 (640ms) | 12 (480ms) | 8 (320ms) |
+|:-----:|:----:|:-----:|:-----:|:-----:|:-----:|:-----:|
+| full  | 2.72%| -     | -     | -     | -     | -     |
+| lc=32 | -    | 3.09% | 3.07% | 3.26% | 3.31% | 3.44% |
+| 16    | -    | 3.10% | 3.07% | 3.27% | 3.32% | 3.50% |
+| 8     | -    | 3.10% | 3.11% | 3.31% | 3.39% | 3.62% |
+| 4     | -    | 3.12% | 3.13% | 3.37% | 3.51% | 3.80% |
+| 2     | -    | 3.19% | 3.24% | 3.50% | 3.79% | 4.38% |
 
 # **About SpeechBrain**
 - Website: https://speechbrain.github.io/

From de7d9971f47f568483da53828857499adb483998 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Thu, 14 Dec 2023 15:05:42 +0100
Subject: [PATCH 42/83] WIP refactor: Introduce DCTConfig dataclass

---
 recipes/LibriSpeech/ASR/transducer/README.md  |   2 +-
 .../hparams/conformer_transducer.yaml         |  20 +-
 recipes/LibriSpeech/ASR/transducer/train.py   |  47 +----
 .../lobes/models/transformer/Branchformer.py  |   3 -
 .../lobes/models/transformer/Conformer.py     |  40 ++--
 .../lobes/models/transformer/Transformer.py   |   2 -
 .../models/transformer/TransformerASR.py      | 198 ++++++++++--------
 speechbrain/utils/DCT.py                      | 130 ++++++++++++
 8 files changed, 275 insertions(+), 167 deletions(-)
 create mode 100644 speechbrain/utils/DCT.py

diff --git a/recipes/LibriSpeech/ASR/transducer/README.md b/recipes/LibriSpeech/ASR/transducer/README.md
index 5a28e6965f..bad54ba73c 100644
--- a/recipes/LibriSpeech/ASR/transducer/README.md
+++ b/recipes/LibriSpeech/ASR/transducer/README.md
@@ -22,7 +22,7 @@ python train.py hparams/conformer_transducer.yaml
 
 # Librispeech Results
 
-Dev. clean is evaluated with Greedy Decoding while the test sets are using Greedy Decoding OR a RNNLM + Beam Search.  
+Dev. clean is evaluated with Greedy Decoding while the test sets are using Greedy Decoding OR a RNNLM + Beam Search.
 Evaluation is performed in fp32.
 
 | Release | Hyperparams file | Train precision | Dev-clean Greedy | Test-clean Greedy | Test-other Greedy | Test-clean BS+RNNLM | Test-other BS+RNNLM | Model link | GPUs |
diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index fd426c853a..69fb96c0b2 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -83,14 +83,18 @@ test_left_context_size: 64
 valid_chunk_size: -1
 valid_left_context_size: -1
 
-dynamic_chunk_prob: 0.6
-dynamic_chunk_min: 8
-dynamic_chunk_max: 32
-
-# Left context gets expressed as a number of chunks
-dynamic_left_context_prob: 0.75
-dynamic_left_context_min: 2
-dynamic_left_context_max: 32
+dct_config_sampler: !new:speechbrain.utils.DCT.DCTConfigRandomSampler # yamllint disable-line rule:line-length
+   dct_prob: 0.6
+
+   chunk_size_min: 8
+   chunk_size_max: 32
+
+   limited_left_context_prob: 0.75
+   left_context_chunks_min: 2
+   left_context_chunks_max: 2
+
+   # valid_config
+   # test_config
 
 # Dataloader options
 train_dataloader_opts:
diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 585a0f622a..a1602e39a0 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -69,60 +69,17 @@ def compute_forward(self, batch, stage):
 
         current_epoch = self.hparams.epoch_counter.current
 
-        # Default to infinite context visibility
-        transformer_chunk_size = -1
-        left_context_chunks = -1
-
         # Old models may not have the streaming hparam, we don't break them in
         # any other way so just check for its presence
         if hasattr(self.hparams, "streaming") and self.hparams.streaming:
-            # TODO: while this is fairly small logic, it may make sense to
-            # extract it to its own class orchestrating dynamic chunk training,
-            # partly because explantions are beneficial
-            if stage == sb.Stage.TRAIN:
-                # When training for streaming, for each batch, we have a
-                # `dynamic_chunk_prob` probability of sampling a chunk size
-                # between `dynamic_chunk_min` and `_max`, otherwise output
-                # frames can see anywhere in the future.
-                # NOTE: We use torch random to be bound to the experiment seed.
-                if torch.rand((1,)).item() < self.hparams.dynamic_chunk_prob:
-                    transformer_chunk_size = torch.randint(
-                        self.hparams.dynamic_chunk_min,
-                        self.hparams.dynamic_chunk_max + 1,
-                        (1,),
-                    ).item()
-
-                # We have a `dynamic_left_context_prob` probability of sampling
-                # a left context size between `dynamic_left_context_min` and
-                # `_max`, otherwise output frames can see anywhere in the past.
-                # Note that this only has an effect when using a chunk size
-                # above.
-                if (
-                    torch.rand((1,)).item()
-                    < self.hparams.dynamic_left_context_prob
-                ):
-                    left_context_chunks = torch.randint(
-                        self.hparams.dynamic_left_context_min,
-                        self.hparams.dynamic_left_context_max + 1,
-                        (1,),
-                    ).item()
-            elif stage == sb.Stage.TEST:
-                transformer_chunk_size = self.hparams.test_chunk_size
-                left_context_chunks = self.hparams.test_left_context_size
-            elif stage == sb.Stage.VALID:
-                transformer_chunk_size = self.hparams.valid_chunk_size
-                left_context_chunks = self.hparams.valid_left_context_size
+            dct_config = self.hparams.dct_config_sampler(stage)
 
         # logger.info(f"Batch uses tfx chunk size = {transformer_chunk_size}, frame chunk_size = {chunk_size}")
         feats = self.modules.normalize(feats, wav_lens, epoch=current_epoch)
 
         src = self.modules.CNN(feats)
         x = self.modules.enc(
-            src,
-            wav_lens,
-            pad_idx=self.hparams.pad_index,
-            chunk_size=transformer_chunk_size,
-            left_context_chunks=left_context_chunks,
+            src, wav_lens, pad_idx=self.hparams.pad_index, dct_config=dct_config
         )
         x = self.modules.proj_enc(x)
 
diff --git a/speechbrain/lobes/models/transformer/Branchformer.py b/speechbrain/lobes/models/transformer/Branchformer.py
index ad03e6a0c2..3f75701a34 100644
--- a/speechbrain/lobes/models/transformer/Branchformer.py
+++ b/speechbrain/lobes/models/transformer/Branchformer.py
@@ -320,7 +320,6 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
-        chunk_size: int = -1,
     ):
         """
         Arguments
@@ -337,8 +336,6 @@ def forward(
             where S is the sequence length, and E is the embedding dimension.
         """
 
-        assert chunk_size == -1, "Encoder does not support streaming"
-
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:
                 raise ValueError(
diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 0ce64f2098..390befb141 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -20,6 +20,7 @@
     MultiheadAttention,
     PositionalwiseFeedForward,
 )
+from speechbrain.utils.DCT import DCTConfig
 from speechbrain.lobes.models.transformer.hypermixing import HyperMixing
 from speechbrain.nnet.normalization import LayerNorm
 from speechbrain.nnet.activations import Swish
@@ -163,7 +164,7 @@ def _do_conv(self, x, inhibit_padding: bool):
 
         return out
 
-    def forward(self, x, mask=None, chunk_size=-1):
+    def forward(self, x, mask=None, dct_config: Optional[DCTConfig] = None):
         """ Processes the input tensor x and returns the output an output tensor"""
 
         # ref: Dynamic chunk convolution for unified streaming and non-streaming
@@ -172,7 +173,7 @@ def forward(self, x, mask=None, chunk_size=-1):
         # split the input into chunks of size `chunk_size`, but for each chunk
         # provide a left context for left chunk dependencies to be possible.
 
-        if chunk_size >= 1:
+        if dct_config is not None:
             # chances are chunking+causal is unintended; i don't know where it
             # may make sense, but if it does to you, feel free to implement it.
             assert (
@@ -182,17 +183,19 @@ def forward(self, x, mask=None, chunk_size=-1):
             batch_size = x.shape[0]
             chunk_left_context = self.padding
 
-            chunk_count = int(math.ceil(x.shape[1] / chunk_size))
+            chunk_count = int(math.ceil(x.shape[1] / dct_config.chunk_size))
 
-            if x.shape[1] % chunk_size != 0:
-                final_right_padding = chunk_size - (x.shape[1] % chunk_size)
+            if x.shape[1] % dct_config.chunk_size != 0:
+                final_right_padding = dct_config.chunk_size - (
+                    x.shape[1] % dct_config.chunk_size
+                )
             else:
                 final_right_padding = 0
 
             # compute the left context that can and should be added, for each
             # chunk. for the first few chunks, we will need to add extra padding
             applied_left_context = [
-                min(chunk_left_context, i * chunk_size,)
+                min(chunk_left_context, i * dct_config.chunk_size,)
                 for i in range(chunk_count)
             ]
 
@@ -202,8 +205,8 @@ def forward(self, x, mask=None, chunk_size=-1):
             out = [
                 x[
                     :,
-                    i * chunk_size
-                    - applied_left_context[i] : (i + 1) * chunk_size,
+                    i * dct_config.chunk_size
+                    - applied_left_context[i] : (i + 1) * dct_config.chunk_size,
                     ...,
                 ]
                 for i in range(chunk_count)
@@ -401,7 +404,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: torch.Tensor = None,
-        chunk_size: int = -1,
+        dct_config: Optional[DCTConfig] = None,
     ):
         """
         Arguments
@@ -414,9 +417,9 @@ def forward(
             The mask for the src keys per batch.
         pos_embs: torch.Tensor, torch.nn.Module, optional
             Module or tensor containing the input sequence positional embeddings
-        chunk_size: int
-            Whether to preform convolution chunking to hide future context,
-            useful for chunked conformers in a dynamic chunk training setting
+        dct_config: Optional[DCTConfig]
+            DCT configuration object for streaming, specifically involved here
+            to apply Dynamic Chunk Convolution to the convolution module.
         """
         conv_mask: Optional[torch.Tensor] = None
         if src_key_padding_mask is not None:
@@ -437,7 +440,7 @@ def forward(
         )
         x = x + skip
         # convolution module
-        x = x + self.convolution_module(x, conv_mask, chunk_size=chunk_size)
+        x = x + self.convolution_module(x, conv_mask, dct_config=dct_config)
         # ffn module
         x = self.norm2(x + 0.5 * self.ffn_module2(x))
         return x, self_attn
@@ -605,7 +608,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
-        chunk_size: int = -1,
+        dct_config: Optional[DCTConfig] = None,
     ):
         """
         Arguments
@@ -620,10 +623,9 @@ def forward(
             Module or tensor containing the input sequence positional embeddings
             If custom pos_embs are given it needs to have the shape (1, 2*S-1, E)
             where S is the sequence length, and E is the embedding dimension.
-        chunk_size: int
-            Whether to preform convolution chunking to hide future context,
-            useful for chunked conformers in a dynamic chunk training setting
-            `-1` ignores chunking
+        dct_config: Optional[DCTConfig]
+            DCT configuration object for streaming, specifically involved here
+            to apply Dynamic Chunk Convolution to the convolution module.
         """
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:
@@ -639,7 +641,7 @@ def forward(
                 src_mask=src_mask,
                 src_key_padding_mask=src_key_padding_mask,
                 pos_embs=pos_embs,
-                chunk_size=chunk_size,
+                dct_config=dct_config,
             )
             attention_lst.append(attention)
         output = self.norm(output)
diff --git a/speechbrain/lobes/models/transformer/Transformer.py b/speechbrain/lobes/models/transformer/Transformer.py
index add6562e7c..012d87ec44 100644
--- a/speechbrain/lobes/models/transformer/Transformer.py
+++ b/speechbrain/lobes/models/transformer/Transformer.py
@@ -526,7 +526,6 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
-        chunk_size: int = -1,
     ):
         """
         Arguments
@@ -538,7 +537,6 @@ def forward(
         src_key_padding_mask : tensor
             The mask for the src keys per batch (optional).
         """
-        assert chunk_size == -1, "Encoder does not support streaming"
         output = src
         if self.layerdrop_prob > 0.0:
             keep_probs = self.rng.random(len(self.layers))
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 6dfdc31eaa..03b6b9857e 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -7,7 +7,7 @@
 from dataclasses import dataclass
 import torch  # noqa 42
 from torch import nn
-from typing import Any, Optional
+from typing import Any, Optional, Tuple
 from speechbrain.nnet.linear import Linear
 from speechbrain.nnet.containers import ModuleList
 from speechbrain.lobes.models.transformer.Transformer import (
@@ -18,14 +18,15 @@
 )
 from speechbrain.nnet.activations import Swish
 from speechbrain.dataio.dataio import length_to_mask
+from speechbrain.utils.DCT import DCTConfig
 
 
 @dataclass
 class TransformerASRStreamingContext:
     """Streaming metadata and state for a `TransformerASR` instance."""
 
-    chunk_size: int
-    """The size of a chunk expressed in input frames to the transformer."""
+    dct_config: DCTConfig
+    """DCT configuration holding chunk size and context size information."""
 
     encoder_context: Any
     """Opaque encoder context information. It is constructed by the encoder's
@@ -34,6 +35,103 @@ class TransformerASRStreamingContext:
     """
 
 
+def make_asr_src_mask(
+    src: torch.Tensor,
+    causal: bool = False,
+    dct_config: Optional[DCTConfig] = None,
+) -> Optional[torch.Tensor]:
+    """Prepare the source mask of shape (TODO)"""
+
+    if causal:
+        assert dct_config is None
+        return get_lookahead_mask(src)
+
+    if dct_config is not None:
+        # init a mask that masks nothing by default
+        # 0 == no mask, 1 == mask
+        src_mask = torch.zeros(
+            (src.shape[1], src.shape[1]), device=src.device, dtype=torch.bool,
+        )
+
+        # The following is not really the sole source used to implement this,
+        # but it helps introduce the concept.
+        # ref: Unified Streaming and Non-streaming Two-pass End-to-end Model
+        # for Speech Recognition
+        # https://arxiv.org/pdf/2012.05481.pdf
+
+        timesteps = src.size(1)
+
+        # mask the future at the right of each chunk
+        for t in range(timesteps):
+            # if we have a chunk size of 8 then:
+            # for 0..7  -> mask 8..
+            # for 8..15 -> mask 16..
+            # etc.
+            next_chunk_index = (t // dct_config.chunk_size) + 1
+            visible_range = next_chunk_index * dct_config.chunk_size
+            src_mask[t, visible_range:] = True
+
+        # mask the past at the left of each chunk (accounting for left context)
+        # only relevant if using left context
+        if not dct_config.is_infinite_left_context():
+            for t in range(timesteps):
+                chunk_index = t // dct_config.chunk_size
+                chunk_first_t = chunk_index * dct_config.chunk_size
+                frame_remaining_context = max(
+                    0, chunk_first_t - dct_config.left_context_chunks
+                )
+
+                # end range is exclusive, so there is no off-by-one here
+                src_mask[t, :frame_remaining_context] = True
+
+        return src_mask
+
+    return None
+
+
+def make_asr_masks(
+    src,
+    tgt=None,
+    wav_len=None,
+    pad_idx=0,
+    causal: bool = False,
+    dct_config: Optional[DCTConfig] = None,
+):
+    """This function generates masks for training the transformer model,
+    opiniated for an ASR context with encoding masks and, optionally, decoding
+    masks (if specifying `tgt`).
+
+    Arguments
+    ---------
+    src : tensor
+        The sequence to the encoder (required).
+    tgt : tensor
+        The sequence to the decoder.
+    pad_idx : int
+        The index for <pad> token (default=0).
+    TODO: some args are missing lol
+    """
+    src_key_padding_mask = None
+
+    # mask out audio beyond the length of audio for each batch
+    if wav_len is not None:
+        abs_len = torch.round(wav_len * src.shape[1])
+        src_key_padding_mask = ~length_to_mask(abs_len).bool()
+
+    # mask out the source
+    src_mask = make_asr_src_mask(src, causal=causal, dct_config=dct_config)
+
+    # If no decoder in the transformer...
+    if tgt is not None:
+        tgt_key_padding_mask = get_key_padding_mask(tgt, pad_idx=pad_idx)
+        tgt_mask = get_lookahead_mask(tgt)
+    else:
+        tgt_key_padding_mask = None
+        tgt_mask = None
+
+    return src_key_padding_mask, tgt_key_padding_mask, src_mask, tgt_mask
+
+
 class TransformerASR(TransformerInterface):
     """This is an implementation of transformer model for ASR.
 
@@ -200,7 +298,7 @@ def forward(self, src, tgt, wav_len=None, pad_idx=0):
             tgt_key_padding_mask,
             src_mask,
             tgt_mask,
-        ) = self.make_masks(src, tgt, wav_len, pad_idx=pad_idx)
+        ) = make_asr_masks(src, tgt, wav_len, pad_idx=pad_idx)
 
         src = self.custom_src_module(src)
         # add pos encoding to queries if are sinusoidal ones else
@@ -246,82 +344,6 @@ def forward(self, src, tgt, wav_len=None, pad_idx=0):
 
         return encoder_out, decoder_out
 
-    def make_masks(
-        self,
-        src,
-        tgt=None,
-        wav_len=None,
-        pad_idx=0,
-        chunk_size: int = -1,
-        left_context_chunks: int = -1,
-    ):
-        """This method generates the masks for training the transformer model.
-
-        Arguments
-        ---------
-        src : tensor
-            The sequence to the encoder (required).
-        tgt : tensor
-            The sequence to the decoder.
-        pad_idx : int
-            The index for <pad> token (default=0).
-        """
-        src_key_padding_mask = None
-        if wav_len is not None:
-            abs_len = torch.round(wav_len * src.shape[1])
-            src_key_padding_mask = ~length_to_mask(abs_len).bool()
-
-        if chunk_size >= 0 or left_context_chunks >= 0:
-            # wav_len unspecified? make a mask that masks nothing by default
-            # 0 == no mask, 1 == mask
-            src_mask = torch.zeros(
-                (src.shape[1], src.shape[1]),
-                device=src.device,
-                dtype=torch.bool,
-            )
-
-        if left_context_chunks >= 0:
-            for i in range(src.shape[1]):
-                if chunk_size >= 0:
-                    current_chunk = (i // chunk_size) * chunk_size
-                    frame_remaining_context = max(
-                        0, current_chunk - left_context_chunks * chunk_size
-                    )
-                else:
-                    frame_remaining_context = 0
-
-                # end range is exclusive, so there is no off-by-one here
-                src_mask[i, :frame_remaining_context] = True
-
-        # The following is not really the sole source used to implement this,
-        # but it introduces the concept.
-        # ref: Unified Streaming and Non-streaming Two-pass End-to-end Model
-        # for Speech Recognition
-        # https://arxiv.org/pdf/2012.05481.pdf
-        if chunk_size >= 0:
-            for i in range(src.shape[1]):
-                # if we have a chunk size of 8 then:
-                # for 0..7  -> mask 8..
-                # for 8..15 -> mask 16..
-                # etc.
-                visible_range = ((i // chunk_size) + 1) * chunk_size
-                src_mask[i, visible_range:] = True
-        else:
-            src_mask = None
-
-        if self.causal:
-            src_mask = get_lookahead_mask(src)
-
-        # If no decoder in the transformer...
-        if tgt is not None:
-            tgt_key_padding_mask = get_key_padding_mask(tgt, pad_idx=pad_idx)
-            tgt_mask = get_lookahead_mask(tgt)
-        else:
-            tgt_key_padding_mask = None
-            tgt_mask = None
-
-        return src_key_padding_mask, tgt_key_padding_mask, src_mask, tgt_mask
-
     @torch.no_grad()
     def decode(self, tgt, encoder_out, enc_len=None):
         """This method implements a decoding step for the transformer model.
@@ -368,8 +390,7 @@ def encode(
         src,
         wav_len=None,
         pad_idx=0,
-        chunk_size=-1,
-        left_context_chunks: int = -1,
+        dct_config: Optional[DCTConfig] = None,
     ):
         """
         Encoder forward pass
@@ -386,13 +407,12 @@ def encode(
             bz, t, ch1, ch2 = src.shape
             src = src.reshape(bz, t, ch1 * ch2)
 
-        (src_key_padding_mask, _, src_mask, _,) = self.make_masks(
+        (src_key_padding_mask, _, src_mask, _,) = make_asr_masks(
             src,
             None,
             wav_len,
             pad_idx=pad_idx,
-            chunk_size=chunk_size,
-            left_context_chunks=left_context_chunks,
+            dct_config=dct_config,
         )
 
         src = self.custom_src_module(src)
@@ -473,14 +493,14 @@ def encode_streaming(self, src, context: TransformerASRStreamingContext):
         )
         return encoder_out
 
-    def make_streaming_context(self, chunk_size: int, encoder_kwargs={}):
+    def make_streaming_context(self, dct_config: DCTConfig, encoder_kwargs={}):
         """Creates a blank streaming context for this transformer and its
         encoder.
 
         Arguments
         ---------
-        chunk_size : int
-            How many frames comprise a chunk.
+        dct_config : DCTConfig
+            Runtime chunkwise attention configuration.
 
         encoder_kwargs : dict
             Parameters to be forward to the encoder's `make_streaming_context`.
@@ -488,7 +508,7 @@ def make_streaming_context(self, chunk_size: int, encoder_kwargs={}):
             encoder.
         """
         return TransformerASRStreamingContext(
-            chunk_size=chunk_size,
+            dct_config=dct_config,
             encoder_context=self.encoder.make_streaming_context(
                 **encoder_kwargs,
             ),
diff --git a/speechbrain/utils/DCT.py b/speechbrain/utils/DCT.py
new file mode 100644
index 0000000000..99cd6de3e2
--- /dev/null
+++ b/speechbrain/utils/DCT.py
@@ -0,0 +1,130 @@
+"""Configuration and utility classes for classes for Dynamic Chunk Training, as
+often used for the training of streaming-capable models in speech recognition.
+
+Authors
+* Sylvain de Langen 2023
+"""
+
+from speechbrain.core import Stage
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+# NOTE: this configuration object is intended to be relatively specific to DCT;
+# if you want to implement a different similar type of chunking different from
+# DCT you should consider using a different object.
+@dataclass
+class DCTConfig:
+    """Dynamic Chunk Training configuration object for use with transformers,
+    often in ASR for streaming.
+
+    This object may be used both to configure masking at training time and for
+    run-time configuration of DCT-ready models."""
+
+    chunk_size: int
+    """Size in frames of a single chunk, always `>0`.
+    If chunkwise streaming should be disabled at some point, pass an optional
+    streaming config parameter."""
+
+    left_context_size: Optional[int]
+    """Number of *chunks* (not frames) visible to the left, always `>=0`.
+    If zero, then chunks can never attend to any past chunk.
+    If `None`, the left context is infinite (but use
+    `.is_fininite_left_context` for such a check)."""
+
+    def is_infinite_left_context(self) -> bool:
+        """Returns true if the left context is infinite (i.e. any chunk can
+        attend to any past frame)."""
+        return self.left_context_size is not None
+
+
+@dataclass
+class DCTConfigRandomSampler:
+    """Helper class to generate a DCTConfig at runtime depending on the current
+    stage."""
+
+    dct_prob: float
+    """When sampling (during `Stage.TRAIN`), the probability that a finite chunk
+    size will be used.
+    In the other case, any chunk can attend to the full past and future
+    context."""
+
+    chunk_size_min: int
+    """When sampling a random chunk size, the minimum chunk size that can be
+    picked."""
+
+    chunk_size_max: int
+    """When sampling a random chunk size, the maximum chunk size that can be
+    picked."""
+
+    limited_left_context_prob: float
+    """When sampling a random chunk size, the probability that the left context
+    will be limited.
+    In the other case, any chunk can attend to the full past context."""
+
+    left_context_chunks_min: int
+    """When sampling a random left context size, the minimum number of left
+    context chunks that can be picked."""
+
+    left_context_chunks_max: int
+    """When sampling a random left context size, the maximum number of left
+    context chunks that can be picked."""
+
+    test_config: Optional[DCTConfig] = None
+    """The configuration that should be used for `Stage.TEST`.
+    When `None`, evaluation is done with full context (i.e. non-streaming)."""
+
+    valid_config: Optional[DCTConfig] = None
+    """The configuration that should be used for `Stage.VALID`.
+    When `None`, evaluation is done with full context (i.e. non-streaming)."""
+
+    def _sample_bool(prob: float) -> bool:
+        """Samples a random boolean with a probability, in a way that depends on
+        PyTorch's RNG seed.
+
+        Arguments
+        ---------
+        prob : float
+            Probability (0..1) to return True (False otherwise)."""
+        return torch.rand((1,)).item() < prob
+
+    def __call__(self, stage: Stage) -> DCTConfig:
+        """Samples a random (or not) DCT configuration depending on the current
+        stage.
+
+        Arguments
+        ---------
+        stage : speechbrain.core.Stage
+            Current stage of training or evaluation.
+            In training mode, a random DCTConfig will be sampled according to
+            the specified probabilities and ranges.
+            In evaluation, the relevant DCTConfig attribute will be picked.
+        """
+        if stage == Stage.TRAIN:
+            # When training for streaming, for each batch, we have a
+            # `dynamic_chunk_prob` probability of sampling a chunk size
+            # between `dynamic_chunk_min` and `_max`, otherwise output
+            # frames can see anywhere in the future.
+            if self._sample_bool(self.dct_prob):
+                chunk_size = torch.randint(
+                    self.chunk_size_min, self.chunk_size_max + 1, (1,),
+                ).item()
+
+                if self._sample_bool(self.limited_left_context_prob):
+                    left_context_chunks = torch.randint(
+                        self.left_context_chunks_min,
+                        self.left_context_chunks_max + 1,
+                        (1,),
+                    ).item()
+                else:
+                    left_context_chunks = None
+
+                return DCTConfig(chunk_size, left_context_chunks)
+            return None
+        elif stage == Stage.TEST:
+            return self.test_config
+        elif stage == Stage.VALID:
+            return self.valid_config
+        else:
+            raise AttributeError(f"Unsupported stage found {stage}")

From d9b6f8886606b295628b652256cdf3740e8ec078 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Thu, 14 Dec 2023 15:35:25 +0100
Subject: [PATCH 43/83] Improved notice in README

---
 recipes/LibriSpeech/ASR/transducer/README.md | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/README.md b/recipes/LibriSpeech/ASR/transducer/README.md
index bad54ba73c..2adfa17865 100644
--- a/recipes/LibriSpeech/ASR/transducer/README.md
+++ b/recipes/LibriSpeech/ASR/transducer/README.md
@@ -20,6 +20,10 @@ pip install numba
 python train.py hparams/conformer_transducer.yaml
 ```
 
+We recommend training with `--precision=bf16` on supported GPUs (e.g. A100) for a significant speedup.  
+`--precision=fp16` was found to be stable for this recipe but might lead to slightly worse model accuracy.  
+Tweaking `max_batch_len` may also prove useful.
+
 # Librispeech Results
 
 Dev. clean is evaluated with Greedy Decoding while the test sets are using Greedy Decoding OR a RNNLM + Beam Search.
@@ -27,7 +31,9 @@ Evaluation is performed in fp32.
 
 | Release | Hyperparams file | Train precision | Dev-clean Greedy | Test-clean Greedy | Test-other Greedy | Test-clean BS+RNNLM | Test-other BS+RNNLM | Model link | GPUs |
 |:-------------:|:---------------------------:|:-:| :------:| :-----------:| :------------------:| :------------------:| :------------------:| :--------:| :-----------:|
-| 2023-12-12 | conformer_transducer.yaml `streaming: True` | bf16 | 2.56% | 2.72% | 6.47% | TBD | TBD | https://drive.google.com/drive/folders/1QtQz1Bkd_QPYnf3CyxhJ57ovbSZC2EhN?usp=sharing | [4x A100SXM4 40GB](https://docs.alliancecan.ca/wiki/Narval/en) |
+| 2023-12-12 | conformer_transducer.yaml `streaming: True` | bf16 | 2.56% | 2.72% | 6.47% | \* | \* | https://drive.google.com/drive/folders/1QtQz1Bkd_QPYnf3CyxhJ57ovbSZC2EhN?usp=sharing | [4x A100SXM4 40GB](https://docs.alliancecan.ca/wiki/Narval/en) |
+
+<sub>\*: not evaluated due to performance issues, see [issue #2301](https://github.com/speechbrain/speechbrain/issues/2301)</sub>
 
 ## Streaming model
 

From 11fec0cd5793b06912e2a38cfe5876a021e2399c Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Thu, 14 Dec 2023 16:08:01 +0100
Subject: [PATCH 44/83] Formatting and linting fixes

---
 .../models/transformer/TransformerASR.py      | 34 ++++++++++++++-----
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 8f16de40da..2256c3faeb 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -7,7 +7,7 @@
 from dataclasses import dataclass
 import torch  # noqa 42
 from torch import nn
-from typing import Any, Optional, Tuple
+from typing import Any, Optional
 from speechbrain.nnet.linear import Linear
 from speechbrain.nnet.containers import ModuleList
 from speechbrain.lobes.models.transformer.Transformer import (
@@ -40,7 +40,24 @@ def make_asr_src_mask(
     causal: bool = False,
     dct_config: Optional[DCTConfig] = None,
 ) -> Optional[torch.Tensor]:
-    """Prepare the source mask of shape (TODO)"""
+    """Prepare the source transformer mask that restricts which frames can
+    attend to which frames depending on causal or other simple restricted
+    attention methods.
+
+    Arguments
+    ---------
+    src: torch.Tensor
+        The source tensor to build a mask from. The contents of the tensor are
+        not actually used currently; only its shape and other metadata (e.g.
+        device).
+
+    causal: bool
+        Whether strict causality shall be used. Frames will not be able to
+        attend to any future frame.
+
+    dct_config: DCTConfig, optional
+        Dynamic Chunk Training configuration. This implements a simple form of
+        chunkwise attention. Incompatible with `causal`. See `DCT`"""
 
     if causal:
         assert dct_config is None
@@ -109,7 +126,10 @@ def make_asr_masks(
         The sequence to the decoder.
     pad_idx : int
         The index for <pad> token (default=0).
-    TODO: some args are missing lol
+    causal: bool
+        Whether strict causality shall be used. See `make_asr_src_mask`
+    dct_config: DCTConfig, optional
+        Dynamic Chunk Training configuration. See `make_asr_src_mask`
     """
     src_key_padding_mask = None
 
@@ -408,11 +428,7 @@ def encode(
             src = src.reshape(bz, t, ch1 * ch2)
 
         (src_key_padding_mask, _, src_mask, _,) = make_asr_masks(
-            src,
-            None,
-            wav_len,
-            pad_idx=pad_idx,
-            dct_config=dct_config,
+            src, None, wav_len, pad_idx=pad_idx, dct_config=dct_config,
         )
 
         src = self.custom_src_module(src)
@@ -429,7 +445,7 @@ def encode(
             src_mask=src_mask,
             src_key_padding_mask=src_key_padding_mask,
             pos_embs=pos_embs_source,
-            chunk_size=chunk_size,
+            dct_config=dct_config,
         )
         return encoder_out
 

From 65255c8fd125779e79cea504d8160c3cdc7572c3 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Thu, 14 Dec 2023 16:36:43 +0100
Subject: [PATCH 45/83] Attempt at fixing circular import?

---
 speechbrain/utils/DCT.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/speechbrain/utils/DCT.py b/speechbrain/utils/DCT.py
index 99cd6de3e2..65949c2ce4 100644
--- a/speechbrain/utils/DCT.py
+++ b/speechbrain/utils/DCT.py
@@ -5,7 +5,7 @@
 * Sylvain de Langen 2023
 """
 
-from speechbrain.core import Stage
+import speechbrain as sb
 from dataclasses import dataclass
 from typing import Optional
 
@@ -89,7 +89,7 @@ def _sample_bool(prob: float) -> bool:
             Probability (0..1) to return True (False otherwise)."""
         return torch.rand((1,)).item() < prob
 
-    def __call__(self, stage: Stage) -> DCTConfig:
+    def __call__(self, stage: "sb.core.Stage") -> DCTConfig:
         """Samples a random (or not) DCT configuration depending on the current
         stage.
 
@@ -101,7 +101,7 @@ def __call__(self, stage: Stage) -> DCTConfig:
             the specified probabilities and ranges.
             In evaluation, the relevant DCTConfig attribute will be picked.
         """
-        if stage == Stage.TRAIN:
+        if stage == sb.core.Stage.TRAIN:
             # When training for streaming, for each batch, we have a
             # `dynamic_chunk_prob` probability of sampling a chunk size
             # between `dynamic_chunk_min` and `_max`, otherwise output
@@ -122,9 +122,9 @@ def __call__(self, stage: Stage) -> DCTConfig:
 
                 return DCTConfig(chunk_size, left_context_chunks)
             return None
-        elif stage == Stage.TEST:
+        elif stage == sb.core.Stage.TEST:
             return self.test_config
-        elif stage == Stage.VALID:
+        elif stage == sb.core.Stage.VALID:
             return self.valid_config
         else:
             raise AttributeError(f"Unsupported stage found {stage}")

From 0ec54175fbc43efcf66aa196c49b3cf89e424ed4 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Thu, 14 Dec 2023 16:41:24 +0100
Subject: [PATCH 46/83] utils can't depend on core it seems; move dct

---
 .../ASR/transducer/hparams/conformer_transducer.yaml            | 2 +-
 speechbrain/{utils => lobes/models/transformer}/DCT.py          | 0
 speechbrain/lobes/models/transformer/TransformerASR.py          | 2 +-
 3 files changed, 2 insertions(+), 2 deletions(-)
 rename speechbrain/{utils => lobes/models/transformer}/DCT.py (100%)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index 69fb96c0b2..c7b1ea4cd3 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -83,7 +83,7 @@ test_left_context_size: 64
 valid_chunk_size: -1
 valid_left_context_size: -1
 
-dct_config_sampler: !new:speechbrain.utils.DCT.DCTConfigRandomSampler # yamllint disable-line rule:line-length
+dct_config_sampler: !new:speechbrain.lobes.models.transformer.DCT.DCTConfigRandomSampler # yamllint disable-line rule:line-length
    dct_prob: 0.6
 
    chunk_size_min: 8
diff --git a/speechbrain/utils/DCT.py b/speechbrain/lobes/models/transformer/DCT.py
similarity index 100%
rename from speechbrain/utils/DCT.py
rename to speechbrain/lobes/models/transformer/DCT.py
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 2256c3faeb..73273f459e 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -18,7 +18,7 @@
 )
 from speechbrain.nnet.activations import Swish
 from speechbrain.dataio.dataio import length_to_mask
-from speechbrain.utils.DCT import DCTConfig
+from speechbrain.lobes.models.transformer.DCT import DCTConfig
 
 
 @dataclass

From 61606ac742355f53186abd3da8e249316b4e27cc Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Thu, 14 Dec 2023 16:45:19 +0100
Subject: [PATCH 47/83] Whoops, missed file

---
 speechbrain/lobes/models/transformer/Conformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 390befb141..42becffdbf 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -20,7 +20,7 @@
     MultiheadAttention,
     PositionalwiseFeedForward,
 )
-from speechbrain.utils.DCT import DCTConfig
+from speechbrain.lobes.models.transformer.DCT import DCTConfig
 from speechbrain.lobes.models.transformer.hypermixing import HyperMixing
 from speechbrain.nnet.normalization import LayerNorm
 from speechbrain.nnet.activations import Swish

From fe38e5b18248154076d2c47c0b969f659ef85925 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 09:47:41 +0100
Subject: [PATCH 48/83] Add DCT test, fix issues

---
 speechbrain/lobes/models/transformer/DCT.py |  6 ++--
 tests/unittests/test_dct.py                 | 36 +++++++++++++++++++++
 2 files changed, 39 insertions(+), 3 deletions(-)
 create mode 100644 tests/unittests/test_dct.py

diff --git a/speechbrain/lobes/models/transformer/DCT.py b/speechbrain/lobes/models/transformer/DCT.py
index 65949c2ce4..3c26bc0115 100644
--- a/speechbrain/lobes/models/transformer/DCT.py
+++ b/speechbrain/lobes/models/transformer/DCT.py
@@ -27,7 +27,7 @@ class DCTConfig:
     If chunkwise streaming should be disabled at some point, pass an optional
     streaming config parameter."""
 
-    left_context_size: Optional[int]
+    left_context_size: Optional[int] = None
     """Number of *chunks* (not frames) visible to the left, always `>=0`.
     If zero, then chunks can never attend to any past chunk.
     If `None`, the left context is infinite (but use
@@ -36,7 +36,7 @@ class DCTConfig:
     def is_infinite_left_context(self) -> bool:
         """Returns true if the left context is infinite (i.e. any chunk can
         attend to any past frame)."""
-        return self.left_context_size is not None
+        return self.left_context_size is None
 
 
 @dataclass
@@ -79,7 +79,7 @@ class DCTConfigRandomSampler:
     """The configuration that should be used for `Stage.VALID`.
     When `None`, evaluation is done with full context (i.e. non-streaming)."""
 
-    def _sample_bool(prob: float) -> bool:
+    def _sample_bool(self, prob: float) -> bool:
         """Samples a random boolean with a probability, in a way that depends on
         PyTorch's RNG seed.
 
diff --git a/tests/unittests/test_dct.py b/tests/unittests/test_dct.py
new file mode 100644
index 0000000000..900f222da1
--- /dev/null
+++ b/tests/unittests/test_dct.py
@@ -0,0 +1,36 @@
+def test_sampler():
+    from speechbrain.core import Stage
+    from speechbrain.lobes.models.transformer.DCT import (
+        DCTConfig, DCTConfigRandomSampler
+    )
+
+    # sanity check and cover for the random smapler
+
+    valid_cfg = DCTConfig(16, 32)
+    test_cfg = DCTConfig(16, 32)
+
+    sampler = DCTConfigRandomSampler(
+        dct_prob=1.0,
+        chunk_size_min=8,
+        chunk_size_max=8,
+
+        limited_left_context_prob=1.0,
+        left_context_chunks_min=16,
+        left_context_chunks_max=16,
+
+        test_config=valid_cfg,
+        valid_config=test_cfg
+    )
+
+    sampled_train_config = sampler(Stage.TRAIN)
+    assert sampled_train_config.chunk_size == 8
+    assert sampled_train_config.left_context_size == 16
+
+    assert(sampler(Stage.VALID) == valid_cfg)
+    assert(sampler(Stage.TEST) == test_cfg)
+
+def test_dct():
+    from speechbrain.lobes.models.transformer.DCT import DCTConfig
+
+    assert DCTConfig(chunk_size=16).is_infinite_left_context()
+    assert not DCTConfig(chunk_size=16, left_context_size=4).is_infinite_left_context()

From c2fc373fd6b34d523911dfd7318888a46cccb443 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 10:20:33 +0100
Subject: [PATCH 49/83] Remove now obsolete yaml variables for streaming

---
 .../ASR/transducer/hparams/conformer_transducer.yaml        | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index c7b1ea4cd3..093381e63d 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -77,12 +77,6 @@ win_length: 32
 # `streaming: False`.
 streaming: True  # controls all DCT & chunk size & left context mechanisms
 
-test_chunk_size: 16
-test_left_context_size: 64
-
-valid_chunk_size: -1
-valid_left_context_size: -1
-
 dct_config_sampler: !new:speechbrain.lobes.models.transformer.DCT.DCTConfigRandomSampler # yamllint disable-line rule:line-length
    dct_prob: 0.6
 

From 2d31242f826e819aaeefaa5bb596ed7bc2ab7bc8 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 10:24:53 +0100
Subject: [PATCH 50/83] Formatting

---
 tests/unittests/test_dct.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/tests/unittests/test_dct.py b/tests/unittests/test_dct.py
index 900f222da1..0a86134d72 100644
--- a/tests/unittests/test_dct.py
+++ b/tests/unittests/test_dct.py
@@ -1,7 +1,8 @@
 def test_sampler():
     from speechbrain.core import Stage
     from speechbrain.lobes.models.transformer.DCT import (
-        DCTConfig, DCTConfigRandomSampler
+        DCTConfig,
+        DCTConfigRandomSampler,
     )
 
     # sanity check and cover for the random smapler
@@ -13,24 +14,25 @@ def test_sampler():
         dct_prob=1.0,
         chunk_size_min=8,
         chunk_size_max=8,
-
         limited_left_context_prob=1.0,
         left_context_chunks_min=16,
         left_context_chunks_max=16,
-
         test_config=valid_cfg,
-        valid_config=test_cfg
+        valid_config=test_cfg,
     )
 
     sampled_train_config = sampler(Stage.TRAIN)
     assert sampled_train_config.chunk_size == 8
     assert sampled_train_config.left_context_size == 16
 
-    assert(sampler(Stage.VALID) == valid_cfg)
-    assert(sampler(Stage.TEST) == test_cfg)
+    assert sampler(Stage.VALID) == valid_cfg
+    assert sampler(Stage.TEST) == test_cfg
+
 
 def test_dct():
     from speechbrain.lobes.models.transformer.DCT import DCTConfig
 
     assert DCTConfig(chunk_size=16).is_infinite_left_context()
-    assert not DCTConfig(chunk_size=16, left_context_size=4).is_infinite_left_context()
+    assert not DCTConfig(
+        chunk_size=16, left_context_size=4
+    ).is_infinite_left_context()

From faa79e92c49249bd4c47131bbe0b426b53afced0 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 10:42:59 +0100
Subject: [PATCH 51/83] Add dummy dct_config parameter to keep unsupported
 encoders working

---
 speechbrain/lobes/models/transformer/Branchformer.py   | 2 ++
 speechbrain/lobes/models/transformer/Transformer.py    | 3 +++
 speechbrain/lobes/models/transformer/TransformerASR.py | 1 +
 3 files changed, 6 insertions(+)

diff --git a/speechbrain/lobes/models/transformer/Branchformer.py b/speechbrain/lobes/models/transformer/Branchformer.py
index 3f75701a34..a3cad26779 100644
--- a/speechbrain/lobes/models/transformer/Branchformer.py
+++ b/speechbrain/lobes/models/transformer/Branchformer.py
@@ -320,6 +320,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
+        dct_config = None,
     ):
         """
         Arguments
@@ -335,6 +336,7 @@ def forward(
             If custom pos_embs are given it needs to have the shape (1, 2*S-1, E)
             where S is the sequence length, and E is the embedding dimension.
         """
+        assert dct_config is None, "DCT unsupported for this encoder"
 
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:
diff --git a/speechbrain/lobes/models/transformer/Transformer.py b/speechbrain/lobes/models/transformer/Transformer.py
index 5b56d4f17b..c51514f0a1 100644
--- a/speechbrain/lobes/models/transformer/Transformer.py
+++ b/speechbrain/lobes/models/transformer/Transformer.py
@@ -530,6 +530,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
+        dct_config = None,
     ):
         """
         Arguments
@@ -541,6 +542,8 @@ def forward(
         src_key_padding_mask : tensor
             The mask for the src keys per batch (optional).
         """
+        assert dct_config is None, "DCT unsupported for this encoder"
+
         output = src
         if self.layerdrop_prob > 0.0:
             keep_probs = self.rng.random(len(self.layers))
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 73273f459e..5bb6210825 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -447,6 +447,7 @@ def encode(
             pos_embs=pos_embs_source,
             dct_config=dct_config,
         )
+
         return encoder_out
 
     def encode_streaming(self, src, context: TransformerASRStreamingContext):

From 90f136702706f927c617a36453319413e26dfa55 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 10:45:44 +0100
Subject: [PATCH 52/83] Linting fix

---
 speechbrain/lobes/models/transformer/Branchformer.py | 2 +-
 speechbrain/lobes/models/transformer/Transformer.py  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Branchformer.py b/speechbrain/lobes/models/transformer/Branchformer.py
index a3cad26779..23cd6dec68 100644
--- a/speechbrain/lobes/models/transformer/Branchformer.py
+++ b/speechbrain/lobes/models/transformer/Branchformer.py
@@ -320,7 +320,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
-        dct_config = None,
+        dct_config=None,
     ):
         """
         Arguments
diff --git a/speechbrain/lobes/models/transformer/Transformer.py b/speechbrain/lobes/models/transformer/Transformer.py
index c51514f0a1..5cf1ca64b8 100644
--- a/speechbrain/lobes/models/transformer/Transformer.py
+++ b/speechbrain/lobes/models/transformer/Transformer.py
@@ -530,7 +530,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
-        dct_config = None,
+        dct_config=None,
     ):
         """
         Arguments

From bcc6b2c23cb942428f75372a3c38a4f02212fa5b Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 11:01:22 +0100
Subject: [PATCH 53/83] Fix typo

---
 speechbrain/lobes/models/transformer/TransformerASR.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 5bb6210825..393613e799 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -95,7 +95,7 @@ def make_asr_src_mask(
                 chunk_index = t // dct_config.chunk_size
                 chunk_first_t = chunk_index * dct_config.chunk_size
                 frame_remaining_context = max(
-                    0, chunk_first_t - dct_config.left_context_chunks
+                    0, chunk_first_t - dct_config.left_context_size
                 )
 
                 # end range is exclusive, so there is no off-by-one here

From 2577cc65e00151222bbf064f3c6e28b57f7a4d99 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 11:14:14 +0100
Subject: [PATCH 54/83] Add note on runtime autocast accuracy

---
 recipes/LibriSpeech/ASR/transducer/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/README.md b/recipes/LibriSpeech/ASR/transducer/README.md
index 387f9c7dd9..0c7706994a 100644
--- a/recipes/LibriSpeech/ASR/transducer/README.md
+++ b/recipes/LibriSpeech/ASR/transducer/README.md
@@ -28,7 +28,7 @@ According to our tests, the performance is not affected.
 # Librispeech Results
 
 Dev. clean is evaluated with Greedy Decoding while the test sets are using Greedy Decoding OR a RNNLM + Beam Search.
-Evaluation is performed in fp32.
+Evaluation is performed in fp32. However, we found that during inference, fp16 or bf16 autocast has very little incidence on the WER.
 
 | Release | Hyperparams file | Train precision | Dev-clean Greedy | Test-clean Greedy | Test-other Greedy | Test-clean BS+RNNLM | Test-other BS+RNNLM | Model link | GPUs |
 |:-------------:|:---------------------------:|:-:| :------:| :-----------:| :------------------:| :------------------:| :------------------:| :--------:| :-----------:|

From 0c8e382950185c2562ef539dec282416a6cb5d9c Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 11:26:06 +0100
Subject: [PATCH 55/83] Fix very bad typo from refactor in YAML

---
 .../ASR/transducer/hparams/conformer_transducer.yaml            | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index 093381e63d..a677a10e2e 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -85,7 +85,7 @@ dct_config_sampler: !new:speechbrain.lobes.models.transformer.DCT.DCTConfigRando
 
    limited_left_context_prob: 0.75
    left_context_chunks_min: 2
-   left_context_chunks_max: 2
+   left_context_chunks_max: 32
 
    # valid_config
    # test_config

From db731140c0233f55bf3881ecf416c8cec0d725d6 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 11:28:53 +0100
Subject: [PATCH 56/83] Fix hasattr streaming check

---
 recipes/LibriSpeech/ASR/transducer/train.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index a1602e39a0..b7d9e792c4 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -73,6 +73,8 @@ def compute_forward(self, batch, stage):
         # any other way so just check for its presence
         if hasattr(self.hparams, "streaming") and self.hparams.streaming:
             dct_config = self.hparams.dct_config_sampler(stage)
+        else:
+            dct_config = None
 
         # logger.info(f"Batch uses tfx chunk size = {transformer_chunk_size}, frame chunk_size = {chunk_size}")
         feats = self.modules.normalize(feats, wav_lens, epoch=current_epoch)

From 74496e611cd1d163a5e3777515130472a9185fc3 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 11:29:27 +0100
Subject: [PATCH 57/83] Remove legacy comment

---
 recipes/LibriSpeech/ASR/transducer/train.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index b7d9e792c4..055b66cef2 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -76,7 +76,6 @@ def compute_forward(self, batch, stage):
         else:
             dct_config = None
 
-        # logger.info(f"Batch uses tfx chunk size = {transformer_chunk_size}, frame chunk_size = {chunk_size}")
         feats = self.modules.normalize(feats, wav_lens, epoch=current_epoch)
 
         src = self.modules.CNN(feats)

From 455823234c94477d921dfa6d08d76ad3ba3bea5e Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 11:40:19 +0100
Subject: [PATCH 58/83] Fix left context size calculation in new mask code

---
 speechbrain/lobes/models/transformer/TransformerASR.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 393613e799..1f0d0e9443 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -94,8 +94,13 @@ def make_asr_src_mask(
             for t in range(timesteps):
                 chunk_index = t // dct_config.chunk_size
                 chunk_first_t = chunk_index * dct_config.chunk_size
+
+                left_context_frames = (
+                    dct_config.left_context_size * dct_config.chunk_size
+                )
+
                 frame_remaining_context = max(
-                    0, chunk_first_t - dct_config.left_context_size
+                    0, chunk_first_t - left_context_frames,
                 )
 
                 # end range is exclusive, so there is no off-by-one here

From 8da79f3141121a60afabd5e644a9decff2038dc0 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 12:04:24 +0100
Subject: [PATCH 59/83] Fix causal models in TransformerASR

---
 .../lobes/models/transformer/TransformerASR.py        | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 1f0d0e9443..b893496790 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -323,7 +323,9 @@ def forward(self, src, tgt, wav_len=None, pad_idx=0):
             tgt_key_padding_mask,
             src_mask,
             tgt_mask,
-        ) = make_asr_masks(src, tgt, wav_len, pad_idx=pad_idx)
+        ) = make_asr_masks(
+            src, tgt, wav_len, causal=self.causal, pad_idx=pad_idx
+        )
 
         src = self.custom_src_module(src)
         # add pos encoding to queries if are sinusoidal ones else
@@ -433,7 +435,12 @@ def encode(
             src = src.reshape(bz, t, ch1 * ch2)
 
         (src_key_padding_mask, _, src_mask, _,) = make_asr_masks(
-            src, None, wav_len, pad_idx=pad_idx, dct_config=dct_config,
+            src,
+            None,
+            wav_len,
+            pad_idx=pad_idx,
+            causal=self.causal,
+            dct_config=dct_config,
         )
 
         src = self.custom_src_module(src)

From bd9f5065f48e35d359c7231d1e6ee74dff8cfdfc Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 15:28:44 +0100
Subject: [PATCH 60/83] Remove comment on high-level inference code

---
 recipes/LibriSpeech/ASR/transducer/README.md | 2 --
 1 file changed, 2 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/README.md b/recipes/LibriSpeech/ASR/transducer/README.md
index 0c7706994a..aa6d964e82 100644
--- a/recipes/LibriSpeech/ASR/transducer/README.md
+++ b/recipes/LibriSpeech/ASR/transducer/README.md
@@ -40,8 +40,6 @@ Evaluation is performed in fp32. However, we found that during inference, fp16 o
 
 ### WER vs chunk size & left context
 
-**Note:** High-level streaming inference code is not currently available.
-
 The following matrix presents the Word Error Rate (WER%) achieved on LibriSpeech
 `test-clean` with various chunk sizes (in ms) and left context sizes (in # of
 chunks).

From 0a49c01d6ae26dbc232041ae852cfde726208e37 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 15:38:15 +0100
Subject: [PATCH 61/83] YAML formatting + commenting dynchunktrain stuff

---
 .../hparams/conformer_transducer.yaml         | 25 +++++++++++--------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index a677a10e2e..1bb77365be 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -77,18 +77,21 @@ win_length: 32
 # `streaming: False`.
 streaming: True  # controls all DCT & chunk size & left context mechanisms
 
+# Configuration for Dynamic Chunk Training.
+# In this model, a chunk is roughly equivalent to 40ms of audio.
 dct_config_sampler: !new:speechbrain.lobes.models.transformer.DCT.DCTConfigRandomSampler # yamllint disable-line rule:line-length
-   dct_prob: 0.6
-
-   chunk_size_min: 8
-   chunk_size_max: 32
-
-   limited_left_context_prob: 0.75
-   left_context_chunks_min: 2
-   left_context_chunks_max: 32
-
-   # valid_config
-   # test_config
+   dct_prob: 0.6 # Probability during a batch to limit attention and sample a random chunk size in the following range
+   chunk_size_min: 8 # Minimum chunk size (if in a DCT batch)
+   chunk_size_max: 32 # Maximum chunk size (if in a DCT batch)
+   limited_left_context_prob: 0.75 # If in a DCT batch, the probability during a batch to restrict left context to a random number of chunks
+   left_context_chunks_min: 2 # Minimum left context size (in # of chunks)
+   left_context_chunks_max: 32 # Maximum left context size (in # of chunks)
+   # If you specify a valid/test config, you can optionally have evaluation be
+   # done with a specific DynChunkTrain configuration.
+   # valid_config: !new:speechbrain.lobes.models.transformer.DCT.DCTConfig
+   #    chunk_size: 24
+   #    left_context_size: 16
+   # test_config: ...
 
 # Dataloader options
 train_dataloader_opts:

From 28cfbb6964d005ad568500268bd8c88769185efd Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 16:12:14 +0100
Subject: [PATCH 62/83] Remove outdated comment about DCConv left contexts

---
 speechbrain/lobes/models/transformer/Conformer.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 42becffdbf..c9f8369b67 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -484,12 +484,6 @@ def forward_streaming(
         # about our chunk's outputs); see above to-do
         x = x[..., -orig_len:, :]
 
-        # TODO: this is slightly suboptimal as this will add left padding inside
-        # the convolution code that we do not need. it would be better to
-        # manually add the right-padding ourselves and disable padding inside
-        # the convolution module for this usecase, but it would need some
-        # refactoring.
-
         if context.dcconv_left_context is not None:
             x = torch.cat((context.dcconv_left_context, x), dim=1)
 

From 246b8b43afe82c8e58bd4fed9f6233c0bd844316 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Fri, 15 Dec 2023 16:14:20 +0100
Subject: [PATCH 63/83] Remove commented out debug prints from TransformerASR

---
 speechbrain/lobes/models/transformer/TransformerASR.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index b893496790..2865634ee1 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -501,11 +501,9 @@ def encode_streaming(self, src, context: TransformerASRStreamingContext):
         # so we craft a dummy empty (uninitialized) tensor to help...
         known_left_context = context.encoder_context.layers[0].mha_left_context
         if known_left_context is None:
-            # print(f"no lc known: using {src.shape}")
             pos_encoding_dummy = src
         else:
             target_shape = list(src.shape)
-            # print(f"computing posemb shape: from {target_shape} with lc {known_left_context.shape}")
             target_shape[-2] += known_left_context.shape[-2]
             pos_encoding_dummy = torch.empty(size=target_shape).to(src)
 

From 49e73ec2b919de8cf491a6b45df9303098af1f3b Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 10:22:15 +0100
Subject: [PATCH 64/83] Move DCT into utils again

---
 .../ASR/transducer/hparams/conformer_transducer.yaml          | 4 ++--
 speechbrain/lobes/models/transformer/Conformer.py             | 2 +-
 speechbrain/lobes/models/transformer/TransformerASR.py        | 2 +-
 .../transformer/DCT.py => utils/dynamic_chunk_training.py}    | 0
 tests/unittests/test_dct.py                                   | 4 ++--
 5 files changed, 6 insertions(+), 6 deletions(-)
 rename speechbrain/{lobes/models/transformer/DCT.py => utils/dynamic_chunk_training.py} (100%)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index 1bb77365be..c98c4c015c 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -79,7 +79,7 @@ streaming: True  # controls all DCT & chunk size & left context mechanisms
 
 # Configuration for Dynamic Chunk Training.
 # In this model, a chunk is roughly equivalent to 40ms of audio.
-dct_config_sampler: !new:speechbrain.lobes.models.transformer.DCT.DCTConfigRandomSampler # yamllint disable-line rule:line-length
+dct_config_sampler: !new:speechbrain.utils.dynamic_chunk_training.DCTConfigRandomSampler # yamllint disable-line rule:line-length
    dct_prob: 0.6 # Probability during a batch to limit attention and sample a random chunk size in the following range
    chunk_size_min: 8 # Minimum chunk size (if in a DCT batch)
    chunk_size_max: 32 # Maximum chunk size (if in a DCT batch)
@@ -88,7 +88,7 @@ dct_config_sampler: !new:speechbrain.lobes.models.transformer.DCT.DCTConfigRando
    left_context_chunks_max: 32 # Maximum left context size (in # of chunks)
    # If you specify a valid/test config, you can optionally have evaluation be
    # done with a specific DynChunkTrain configuration.
-   # valid_config: !new:speechbrain.lobes.models.transformer.DCT.DCTConfig
+   # valid_config: !new:speechbrain.utils.dynamic_chunk_training.DCTConfig
    #    chunk_size: 24
    #    left_context_size: 16
    # test_config: ...
diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index c9f8369b67..0c52b3a49a 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -20,7 +20,7 @@
     MultiheadAttention,
     PositionalwiseFeedForward,
 )
-from speechbrain.lobes.models.transformer.DCT import DCTConfig
+from speechbrain.utils.dynamic_chunk_training import DCTConfig
 from speechbrain.lobes.models.transformer.hypermixing import HyperMixing
 from speechbrain.nnet.normalization import LayerNorm
 from speechbrain.nnet.activations import Swish
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 2865634ee1..8e473a28a0 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -18,7 +18,7 @@
 )
 from speechbrain.nnet.activations import Swish
 from speechbrain.dataio.dataio import length_to_mask
-from speechbrain.lobes.models.transformer.DCT import DCTConfig
+from speechbrain.utils.dynamic_chunk_training import DCTConfig
 
 
 @dataclass
diff --git a/speechbrain/lobes/models/transformer/DCT.py b/speechbrain/utils/dynamic_chunk_training.py
similarity index 100%
rename from speechbrain/lobes/models/transformer/DCT.py
rename to speechbrain/utils/dynamic_chunk_training.py
diff --git a/tests/unittests/test_dct.py b/tests/unittests/test_dct.py
index 0a86134d72..77d9891fa5 100644
--- a/tests/unittests/test_dct.py
+++ b/tests/unittests/test_dct.py
@@ -1,6 +1,6 @@
 def test_sampler():
     from speechbrain.core import Stage
-    from speechbrain.lobes.models.transformer.DCT import (
+    from speechbrain.utils.dynamic_chunk_training import (
         DCTConfig,
         DCTConfigRandomSampler,
     )
@@ -30,7 +30,7 @@ def test_sampler():
 
 
 def test_dct():
-    from speechbrain.lobes.models.transformer.DCT import DCTConfig
+    from speechbrain.utils.dynamic_chunk_training import DCTConfig
 
     assert DCTConfig(chunk_size=16).is_infinite_left_context()
     assert not DCTConfig(

From 2faf306e89f815421409ef93eeb57532b8cad76e Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 10:46:55 +0100
Subject: [PATCH 65/83] Rename all(?) mentions of DCT to explicit dynamic chunk
 training

---
 .../hparams/conformer_transducer.yaml         | 14 ++---
 recipes/LibriSpeech/ASR/transducer/train.py   | 11 ++--
 .../lobes/models/transformer/Branchformer.py  |  6 ++-
 .../lobes/models/transformer/Conformer.py     | 52 ++++++++++++-------
 .../lobes/models/transformer/Transformer.py   |  6 ++-
 .../models/transformer/TransformerASR.py      | 52 +++++++++++--------
 speechbrain/utils/dynamic_chunk_training.py   | 37 ++++++-------
 ..._dct.py => test_dynamic_chunk_training.py} | 22 ++++----
 8 files changed, 114 insertions(+), 86 deletions(-)
 rename tests/unittests/{test_dct.py => test_dynamic_chunk_training.py} (59%)

diff --git a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
index c98c4c015c..c7ad99c638 100644
--- a/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
+++ b/recipes/LibriSpeech/ASR/transducer/hparams/conformer_transducer.yaml
@@ -75,20 +75,20 @@ win_length: 32
 # At least for the current architecture on LibriSpeech, we found out that
 # non-streaming accuracy is very similar between `streaming: True` and
 # `streaming: False`.
-streaming: True  # controls all DCT & chunk size & left context mechanisms
+streaming: True  # controls all Dynamic Chunk Training & chunk size & left context mechanisms
 
 # Configuration for Dynamic Chunk Training.
 # In this model, a chunk is roughly equivalent to 40ms of audio.
-dct_config_sampler: !new:speechbrain.utils.dynamic_chunk_training.DCTConfigRandomSampler # yamllint disable-line rule:line-length
-   dct_prob: 0.6 # Probability during a batch to limit attention and sample a random chunk size in the following range
-   chunk_size_min: 8 # Minimum chunk size (if in a DCT batch)
-   chunk_size_max: 32 # Maximum chunk size (if in a DCT batch)
-   limited_left_context_prob: 0.75 # If in a DCT batch, the probability during a batch to restrict left context to a random number of chunks
+dynchunktrain_config_sampler: !new:speechbrain.utils.dynamic_chunk_training.DynChunkTrainConfigRandomSampler # yamllint disable-line rule:line-length
+   chunkwise_prob: 0.6 # Probability during a batch to limit attention and sample a random chunk size in the following range
+   chunk_size_min: 8 # Minimum chunk size (if in a DynChunkTrain batch)
+   chunk_size_max: 32 # Maximum chunk size (if in a DynChunkTrain batch)
+   limited_left_context_prob: 0.75 # If in a DynChunkTrain batch, the probability during a batch to restrict left context to a random number of chunks
    left_context_chunks_min: 2 # Minimum left context size (in # of chunks)
    left_context_chunks_max: 32 # Maximum left context size (in # of chunks)
    # If you specify a valid/test config, you can optionally have evaluation be
    # done with a specific DynChunkTrain configuration.
-   # valid_config: !new:speechbrain.utils.dynamic_chunk_training.DCTConfig
+   # valid_config: !new:speechbrain.utils.dynamic_chunk_training.DynChunkTrainConfig
    #    chunk_size: 24
    #    left_context_size: 16
    # test_config: ...
diff --git a/recipes/LibriSpeech/ASR/transducer/train.py b/recipes/LibriSpeech/ASR/transducer/train.py
index 055b66cef2..4a0e889a41 100644
--- a/recipes/LibriSpeech/ASR/transducer/train.py
+++ b/recipes/LibriSpeech/ASR/transducer/train.py
@@ -72,15 +72,20 @@ def compute_forward(self, batch, stage):
         # Old models may not have the streaming hparam, we don't break them in
         # any other way so just check for its presence
         if hasattr(self.hparams, "streaming") and self.hparams.streaming:
-            dct_config = self.hparams.dct_config_sampler(stage)
+            dynchunktrain_config = self.hparams.dynchunktrain_config_sampler(
+                stage
+            )
         else:
-            dct_config = None
+            dynchunktrain_config = None
 
         feats = self.modules.normalize(feats, wav_lens, epoch=current_epoch)
 
         src = self.modules.CNN(feats)
         x = self.modules.enc(
-            src, wav_lens, pad_idx=self.hparams.pad_index, dct_config=dct_config
+            src,
+            wav_lens,
+            pad_idx=self.hparams.pad_index,
+            dynchunktrain_config=dynchunktrain_config,
         )
         x = self.modules.proj_enc(x)
 
diff --git a/speechbrain/lobes/models/transformer/Branchformer.py b/speechbrain/lobes/models/transformer/Branchformer.py
index 23cd6dec68..c2dd297fc7 100644
--- a/speechbrain/lobes/models/transformer/Branchformer.py
+++ b/speechbrain/lobes/models/transformer/Branchformer.py
@@ -320,7 +320,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
-        dct_config=None,
+        dynchunktrain_config=None,
     ):
         """
         Arguments
@@ -336,7 +336,9 @@ def forward(
             If custom pos_embs are given it needs to have the shape (1, 2*S-1, E)
             where S is the sequence length, and E is the embedding dimension.
         """
-        assert dct_config is None, "DCT unsupported for this encoder"
+        assert (
+            dynchunktrain_config is None
+        ), "Dynamic Chunk Training unsupported for this encoder"
 
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:
diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 0c52b3a49a..10259e4562 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -20,7 +20,7 @@
     MultiheadAttention,
     PositionalwiseFeedForward,
 )
-from speechbrain.utils.dynamic_chunk_training import DCTConfig
+from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
 from speechbrain.lobes.models.transformer.hypermixing import HyperMixing
 from speechbrain.nnet.normalization import LayerNorm
 from speechbrain.nnet.activations import Swish
@@ -164,7 +164,12 @@ def _do_conv(self, x, inhibit_padding: bool):
 
         return out
 
-    def forward(self, x, mask=None, dct_config: Optional[DCTConfig] = None):
+    def forward(
+        self,
+        x,
+        mask=None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
+    ):
         """ Processes the input tensor x and returns the output an output tensor"""
 
         # ref: Dynamic chunk convolution for unified streaming and non-streaming
@@ -173,7 +178,7 @@ def forward(self, x, mask=None, dct_config: Optional[DCTConfig] = None):
         # split the input into chunks of size `chunk_size`, but for each chunk
         # provide a left context for left chunk dependencies to be possible.
 
-        if dct_config is not None:
+        if dynchunktrain_config is not None:
             # chances are chunking+causal is unintended; i don't know where it
             # may make sense, but if it does to you, feel free to implement it.
             assert (
@@ -183,11 +188,13 @@ def forward(self, x, mask=None, dct_config: Optional[DCTConfig] = None):
             batch_size = x.shape[0]
             chunk_left_context = self.padding
 
-            chunk_count = int(math.ceil(x.shape[1] / dct_config.chunk_size))
+            chunk_count = int(
+                math.ceil(x.shape[1] / dynchunktrain_config.chunk_size)
+            )
 
-            if x.shape[1] % dct_config.chunk_size != 0:
-                final_right_padding = dct_config.chunk_size - (
-                    x.shape[1] % dct_config.chunk_size
+            if x.shape[1] % dynchunktrain_config.chunk_size != 0:
+                final_right_padding = dynchunktrain_config.chunk_size - (
+                    x.shape[1] % dynchunktrain_config.chunk_size
                 )
             else:
                 final_right_padding = 0
@@ -195,7 +202,7 @@ def forward(self, x, mask=None, dct_config: Optional[DCTConfig] = None):
             # compute the left context that can and should be added, for each
             # chunk. for the first few chunks, we will need to add extra padding
             applied_left_context = [
-                min(chunk_left_context, i * dct_config.chunk_size,)
+                min(chunk_left_context, i * dynchunktrain_config.chunk_size,)
                 for i in range(chunk_count)
             ]
 
@@ -205,8 +212,9 @@ def forward(self, x, mask=None, dct_config: Optional[DCTConfig] = None):
             out = [
                 x[
                     :,
-                    i * dct_config.chunk_size
-                    - applied_left_context[i] : (i + 1) * dct_config.chunk_size,
+                    i * dynchunktrain_config.chunk_size
+                    - applied_left_context[i] : (i + 1)
+                    * dynchunktrain_config.chunk_size,
                     ...,
                 ]
                 for i in range(chunk_count)
@@ -404,7 +412,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: torch.Tensor = None,
-        dct_config: Optional[DCTConfig] = None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
     ):
         """
         Arguments
@@ -417,9 +425,10 @@ def forward(
             The mask for the src keys per batch.
         pos_embs: torch.Tensor, torch.nn.Module, optional
             Module or tensor containing the input sequence positional embeddings
-        dct_config: Optional[DCTConfig]
-            DCT configuration object for streaming, specifically involved here
-            to apply Dynamic Chunk Convolution to the convolution module.
+        dynchunktrain_config: Optional[DynChunkTrainConfig]
+            Dynamic Chunk Training configuration object for streaming,
+            specifically involved here to apply Dynamic Chunk Convolution to
+            the convolution module.
         """
         conv_mask: Optional[torch.Tensor] = None
         if src_key_padding_mask is not None:
@@ -440,7 +449,9 @@ def forward(
         )
         x = x + skip
         # convolution module
-        x = x + self.convolution_module(x, conv_mask, dct_config=dct_config)
+        x = x + self.convolution_module(
+            x, conv_mask, dynchunktrain_config=dynchunktrain_config
+        )
         # ffn module
         x = self.norm2(x + 0.5 * self.ffn_module2(x))
         return x, self_attn
@@ -602,7 +613,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
-        dct_config: Optional[DCTConfig] = None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
     ):
         """
         Arguments
@@ -617,9 +628,10 @@ def forward(
             Module or tensor containing the input sequence positional embeddings
             If custom pos_embs are given it needs to have the shape (1, 2*S-1, E)
             where S is the sequence length, and E is the embedding dimension.
-        dct_config: Optional[DCTConfig]
-            DCT configuration object for streaming, specifically involved here
-            to apply Dynamic Chunk Convolution to the convolution module.
+        dynchunktrain_config: Optional[DynChunkTrainConfig]
+            Dynamic Chunk Training configuration object for streaming,
+            specifically involved here to apply Dynamic Chunk Convolution to the
+            convolution module.
         """
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:
@@ -635,7 +647,7 @@ def forward(
                 src_mask=src_mask,
                 src_key_padding_mask=src_key_padding_mask,
                 pos_embs=pos_embs,
-                dct_config=dct_config,
+                dynchunktrain_config=dynchunktrain_config,
             )
             attention_lst.append(attention)
         output = self.norm(output)
diff --git a/speechbrain/lobes/models/transformer/Transformer.py b/speechbrain/lobes/models/transformer/Transformer.py
index 5cf1ca64b8..52bb367cdb 100644
--- a/speechbrain/lobes/models/transformer/Transformer.py
+++ b/speechbrain/lobes/models/transformer/Transformer.py
@@ -530,7 +530,7 @@ def forward(
         src_mask: Optional[torch.Tensor] = None,
         src_key_padding_mask: Optional[torch.Tensor] = None,
         pos_embs: Optional[torch.Tensor] = None,
-        dct_config=None,
+        dynchunktrain_config=None,
     ):
         """
         Arguments
@@ -542,7 +542,9 @@ def forward(
         src_key_padding_mask : tensor
             The mask for the src keys per batch (optional).
         """
-        assert dct_config is None, "DCT unsupported for this encoder"
+        assert (
+            dynchunktrain_config is None
+        ), "Dynamic Chunk Training unsupported for this encoder"
 
         output = src
         if self.layerdrop_prob > 0.0:
diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 8e473a28a0..dfc16c58ac 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -18,15 +18,16 @@
 )
 from speechbrain.nnet.activations import Swish
 from speechbrain.dataio.dataio import length_to_mask
-from speechbrain.utils.dynamic_chunk_training import DCTConfig
+from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
 
 
 @dataclass
 class TransformerASRStreamingContext:
     """Streaming metadata and state for a `TransformerASR` instance."""
 
-    dct_config: DCTConfig
-    """DCT configuration holding chunk size and context size information."""
+    dynchunktrain_config: DynChunkTrainConfig
+    """Dynamic Chunk Training configuration holding chunk size and context size
+    information."""
 
     encoder_context: Any
     """Opaque encoder context information. It is constructed by the encoder's
@@ -38,7 +39,7 @@ class TransformerASRStreamingContext:
 def make_asr_src_mask(
     src: torch.Tensor,
     causal: bool = False,
-    dct_config: Optional[DCTConfig] = None,
+    dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
 ) -> Optional[torch.Tensor]:
     """Prepare the source transformer mask that restricts which frames can
     attend to which frames depending on causal or other simple restricted
@@ -55,15 +56,15 @@ def make_asr_src_mask(
         Whether strict causality shall be used. Frames will not be able to
         attend to any future frame.
 
-    dct_config: DCTConfig, optional
+    dynchunktrain_config: DynChunkTrainConfig, optional
         Dynamic Chunk Training configuration. This implements a simple form of
-        chunkwise attention. Incompatible with `causal`. See `DCT`"""
+        chunkwise attention. Incompatible with `causal`."""
 
     if causal:
-        assert dct_config is None
+        assert dynchunktrain_config is None
         return get_lookahead_mask(src)
 
-    if dct_config is not None:
+    if dynchunktrain_config is not None:
         # init a mask that masks nothing by default
         # 0 == no mask, 1 == mask
         src_mask = torch.zeros(
@@ -84,19 +85,20 @@ def make_asr_src_mask(
             # for 0..7  -> mask 8..
             # for 8..15 -> mask 16..
             # etc.
-            next_chunk_index = (t // dct_config.chunk_size) + 1
-            visible_range = next_chunk_index * dct_config.chunk_size
+            next_chunk_index = (t // dynchunktrain_config.chunk_size) + 1
+            visible_range = next_chunk_index * dynchunktrain_config.chunk_size
             src_mask[t, visible_range:] = True
 
         # mask the past at the left of each chunk (accounting for left context)
         # only relevant if using left context
-        if not dct_config.is_infinite_left_context():
+        if not dynchunktrain_config.is_infinite_left_context():
             for t in range(timesteps):
-                chunk_index = t // dct_config.chunk_size
-                chunk_first_t = chunk_index * dct_config.chunk_size
+                chunk_index = t // dynchunktrain_config.chunk_size
+                chunk_first_t = chunk_index * dynchunktrain_config.chunk_size
 
                 left_context_frames = (
-                    dct_config.left_context_size * dct_config.chunk_size
+                    dynchunktrain_config.left_context_size
+                    * dynchunktrain_config.chunk_size
                 )
 
                 frame_remaining_context = max(
@@ -117,7 +119,7 @@ def make_asr_masks(
     wav_len=None,
     pad_idx=0,
     causal: bool = False,
-    dct_config: Optional[DCTConfig] = None,
+    dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
 ):
     """This function generates masks for training the transformer model,
     opiniated for an ASR context with encoding masks and, optionally, decoding
@@ -133,7 +135,7 @@ def make_asr_masks(
         The index for <pad> token (default=0).
     causal: bool
         Whether strict causality shall be used. See `make_asr_src_mask`
-    dct_config: DCTConfig, optional
+    dynchunktrain_config: DynChunkTrainConfig, optional
         Dynamic Chunk Training configuration. See `make_asr_src_mask`
     """
     src_key_padding_mask = None
@@ -144,7 +146,9 @@ def make_asr_masks(
         src_key_padding_mask = ~length_to_mask(abs_len).bool()
 
     # mask out the source
-    src_mask = make_asr_src_mask(src, causal=causal, dct_config=dct_config)
+    src_mask = make_asr_src_mask(
+        src, causal=causal, dynchunktrain_config=dynchunktrain_config
+    )
 
     # If no decoder in the transformer...
     if tgt is not None:
@@ -417,7 +421,7 @@ def encode(
         src,
         wav_len=None,
         pad_idx=0,
-        dct_config: Optional[DCTConfig] = None,
+        dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
     ):
         """
         Encoder forward pass
@@ -440,7 +444,7 @@ def encode(
             wav_len,
             pad_idx=pad_idx,
             causal=self.causal,
-            dct_config=dct_config,
+            dynchunktrain_config=dynchunktrain_config,
         )
 
         src = self.custom_src_module(src)
@@ -457,7 +461,7 @@ def encode(
             src_mask=src_mask,
             src_key_padding_mask=src_key_padding_mask,
             pos_embs=pos_embs_source,
-            dct_config=dct_config,
+            dynchunktrain_config=dynchunktrain_config,
         )
 
         return encoder_out
@@ -520,13 +524,15 @@ def encode_streaming(self, src, context: TransformerASRStreamingContext):
         )
         return encoder_out
 
-    def make_streaming_context(self, dct_config: DCTConfig, encoder_kwargs={}):
+    def make_streaming_context(
+        self, dynchunktrain_config: DynChunkTrainConfig, encoder_kwargs={}
+    ):
         """Creates a blank streaming context for this transformer and its
         encoder.
 
         Arguments
         ---------
-        dct_config : DCTConfig
+        dynchunktrain_config : DynChunkTrainConfig
             Runtime chunkwise attention configuration.
 
         encoder_kwargs : dict
@@ -535,7 +541,7 @@ def make_streaming_context(self, dct_config: DCTConfig, encoder_kwargs={}):
             encoder.
         """
         return TransformerASRStreamingContext(
-            dct_config=dct_config,
+            dynchunktrain_config=dynchunktrain_config,
             encoder_context=self.encoder.make_streaming_context(
                 **encoder_kwargs,
             ),
diff --git a/speechbrain/utils/dynamic_chunk_training.py b/speechbrain/utils/dynamic_chunk_training.py
index 3c26bc0115..be2c196961 100644
--- a/speechbrain/utils/dynamic_chunk_training.py
+++ b/speechbrain/utils/dynamic_chunk_training.py
@@ -11,16 +11,16 @@
 
 import torch
 
-# NOTE: this configuration object is intended to be relatively specific to DCT;
-# if you want to implement a different similar type of chunking different from
-# DCT you should consider using a different object.
+# NOTE: this configuration object is intended to be relatively specific to
+# Dynamic Chunk Training; if you want to implement a different similar type of
+# chunking different from that, you should consider using a different object.
 @dataclass
-class DCTConfig:
+class DynChunkTrainConfig:
     """Dynamic Chunk Training configuration object for use with transformers,
     often in ASR for streaming.
 
     This object may be used both to configure masking at training time and for
-    run-time configuration of DCT-ready models."""
+    run-time configuration of DynChunkTrain-ready models."""
 
     chunk_size: int
     """Size in frames of a single chunk, always `>0`.
@@ -40,11 +40,11 @@ def is_infinite_left_context(self) -> bool:
 
 
 @dataclass
-class DCTConfigRandomSampler:
-    """Helper class to generate a DCTConfig at runtime depending on the current
+class DynChunkTrainConfigRandomSampler:
+    """Helper class to generate a DynChunkTrainConfig at runtime depending on the current
     stage."""
 
-    dct_prob: float
+    chunkwise_prob: float
     """When sampling (during `Stage.TRAIN`), the probability that a finite chunk
     size will be used.
     In the other case, any chunk can attend to the full past and future
@@ -71,11 +71,11 @@ class DCTConfigRandomSampler:
     """When sampling a random left context size, the maximum number of left
     context chunks that can be picked."""
 
-    test_config: Optional[DCTConfig] = None
+    test_config: Optional[DynChunkTrainConfig] = None
     """The configuration that should be used for `Stage.TEST`.
     When `None`, evaluation is done with full context (i.e. non-streaming)."""
 
-    valid_config: Optional[DCTConfig] = None
+    valid_config: Optional[DynChunkTrainConfig] = None
     """The configuration that should be used for `Stage.VALID`.
     When `None`, evaluation is done with full context (i.e. non-streaming)."""
 
@@ -89,24 +89,25 @@ def _sample_bool(self, prob: float) -> bool:
             Probability (0..1) to return True (False otherwise)."""
         return torch.rand((1,)).item() < prob
 
-    def __call__(self, stage: "sb.core.Stage") -> DCTConfig:
-        """Samples a random (or not) DCT configuration depending on the current
-        stage.
+    def __call__(self, stage: "sb.core.Stage") -> DynChunkTrainConfig:
+        """In training stage, samples a random DynChunkTrain configuration.
+        During validation or testing, returns the relevant configuration.
 
         Arguments
         ---------
         stage : speechbrain.core.Stage
             Current stage of training or evaluation.
-            In training mode, a random DCTConfig will be sampled according to
-            the specified probabilities and ranges.
-            In evaluation, the relevant DCTConfig attribute will be picked.
+            In training mode, a random DynChunkTrainConfig will be sampled
+            according to the specified probabilities and ranges.
+            During evaluation, the relevant DynChunkTrainConfig attribute will
+            be picked.
         """
         if stage == sb.core.Stage.TRAIN:
             # When training for streaming, for each batch, we have a
             # `dynamic_chunk_prob` probability of sampling a chunk size
             # between `dynamic_chunk_min` and `_max`, otherwise output
             # frames can see anywhere in the future.
-            if self._sample_bool(self.dct_prob):
+            if self._sample_bool(self.chunkwise_prob):
                 chunk_size = torch.randint(
                     self.chunk_size_min, self.chunk_size_max + 1, (1,),
                 ).item()
@@ -120,7 +121,7 @@ def __call__(self, stage: "sb.core.Stage") -> DCTConfig:
                 else:
                     left_context_chunks = None
 
-                return DCTConfig(chunk_size, left_context_chunks)
+                return DynChunkTrainConfig(chunk_size, left_context_chunks)
             return None
         elif stage == sb.core.Stage.TEST:
             return self.test_config
diff --git a/tests/unittests/test_dct.py b/tests/unittests/test_dynamic_chunk_training.py
similarity index 59%
rename from tests/unittests/test_dct.py
rename to tests/unittests/test_dynamic_chunk_training.py
index 77d9891fa5..6c74d10d1b 100644
--- a/tests/unittests/test_dct.py
+++ b/tests/unittests/test_dynamic_chunk_training.py
@@ -1,17 +1,17 @@
-def test_sampler():
+def test_dynchunktrain_sampler():
     from speechbrain.core import Stage
     from speechbrain.utils.dynamic_chunk_training import (
-        DCTConfig,
-        DCTConfigRandomSampler,
+        DynChunkTrainConfig,
+        DynChunkTrainConfigRandomSampler,
     )
 
     # sanity check and cover for the random smapler
 
-    valid_cfg = DCTConfig(16, 32)
-    test_cfg = DCTConfig(16, 32)
+    valid_cfg = DynChunkTrainConfig(16, 32)
+    test_cfg = DynChunkTrainConfig(16, 32)
 
-    sampler = DCTConfigRandomSampler(
-        dct_prob=1.0,
+    sampler = DynChunkTrainConfigRandomSampler(
+        chunkwise_prob=1.0,
         chunk_size_min=8,
         chunk_size_max=8,
         limited_left_context_prob=1.0,
@@ -29,10 +29,10 @@ def test_sampler():
     assert sampler(Stage.TEST) == test_cfg
 
 
-def test_dct():
-    from speechbrain.utils.dynamic_chunk_training import DCTConfig
+def test_dynchunktrain():
+    from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
 
-    assert DCTConfig(chunk_size=16).is_infinite_left_context()
-    assert not DCTConfig(
+    assert DynChunkTrainConfig(chunk_size=16).is_infinite_left_context()
+    assert not DynChunkTrainConfig(
         chunk_size=16, left_context_size=4
     ).is_infinite_left_context()

From c577c5bfb9ebb2b98b7750f45c62806567c4f875 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 11:17:04 +0100
Subject: [PATCH 66/83] Clarify padding logic

---
 speechbrain/lobes/models/transformer/Conformer.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 10259e4562..acb9c55dea 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -227,11 +227,14 @@ def forward(
 
             # TODO: experiment around reflect padding, which is difficult
             # because small chunks have too little time steps to reflect from
+
+            # pad zeroes manually along the time axis
             out = [
                 F.pad(
                     out[i],
                     (
-                        # channel dims, we do not to pad these
+                        # last channel is the channel dim, so do not insert any
+                        # padding at the start or end of that dimension
                         0,
                         0,
                         # add missing left 0-padding if we lacked left context

From 176f1d8645842e8429b52bb5d69d67816a0f7eb9 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 11:25:30 +0100
Subject: [PATCH 67/83] Remove now-useless _do_conv, fix horrible formatting

---
 .../lobes/models/transformer/Conformer.py     | 48 ++++++++-----------
 1 file changed, 21 insertions(+), 27 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index acb9c55dea..3495edd16d 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -142,28 +142,6 @@ def __init__(
             nn.Dropout(dropout),
         )
 
-    def _do_conv(self, x, inhibit_padding: bool):
-        if not inhibit_padding:
-            out = self.conv(x)
-        else:
-            # let's keep backwards compat by pointing at the weights from the
-            # already declared Conv1d.
-
-            # we do not need to edit the bottleneck as it is pointwise (i.e.
-            # time step by time step), thus, it doesn't need padding along the
-            # time dimension
-            out = F.conv1d(
-                x,
-                weight=self.conv.weight,
-                bias=self.conv.bias,
-                stride=self.conv.stride,
-                padding=0,
-                dilation=self.conv.dilation,
-                groups=self.conv.groups,
-            )
-
-        return out
-
     def forward(
         self,
         x,
@@ -212,9 +190,10 @@ def forward(
             out = [
                 x[
                     :,
-                    i * dynchunktrain_config.chunk_size
-                    - applied_left_context[i] : (i + 1)
-                    * dynchunktrain_config.chunk_size,
+                    (
+                        i * dynchunktrain_config.chunk_size
+                        - applied_left_context[i]
+                    ) : ((i + 1) * dynchunktrain_config.chunk_size),
                     ...,
                 ]
                 for i in range(chunk_count)
@@ -263,8 +242,23 @@ def forward(
             # -> [batch_size * num_chunks, in_channels, chunk_size + lc + rpad]
             out = out.transpose(1, 2)
 
+            # let's keep backwards compat by pointing at the weights from the
+            # already declared Conv1d.
+
+            # we do not need to edit the bottleneck as it is pointwise (i.e.
+            # time step by time step), thus, it doesn't need padding along the
+            # time dimension
+
             # -> [batch_size * num_chunks, out_channels, chunk_size + rpad]
-            out = self._do_conv(out, inhibit_padding=True)
+            out = F.conv1d(
+                x,
+                weight=self.conv.weight,
+                bias=self.conv.bias,
+                stride=self.conv.stride,
+                padding=0,
+                dilation=self.conv.dilation,
+                groups=self.conv.groups,
+            )
 
             # -> [batch_size * num_chunks, chunk_size + rpad, out_channels]
             out = out.transpose(1, 2)
@@ -284,7 +278,7 @@ def forward(
             out = self.layer_norm(x)
             out = out.transpose(1, 2)
             out = self.bottleneck(out)
-            out = self._do_conv(out, inhibit_padding=False)
+            out = self.conv(out)
 
             if self.causal:
                 # chomp

From f17470edea2ae990c80f11953ee7030c282a9ea9 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 11:31:23 +0100
Subject: [PATCH 68/83] Slightly fix formatting further

---
 .../lobes/models/transformer/Conformer.py     | 20 ++++++++-----------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 3495edd16d..96dd64eca5 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -163,24 +163,21 @@ def forward(
                 not self.causal
             ), "Chunked convolution not supported with causal padding"
 
+            chunk_size = dynchunktrain_config.chunk_size
             batch_size = x.shape[0]
             chunk_left_context = self.padding
 
-            chunk_count = int(
-                math.ceil(x.shape[1] / dynchunktrain_config.chunk_size)
-            )
+            chunk_count = int(math.ceil(x.shape[1] / chunk_size))
 
-            if x.shape[1] % dynchunktrain_config.chunk_size != 0:
-                final_right_padding = dynchunktrain_config.chunk_size - (
-                    x.shape[1] % dynchunktrain_config.chunk_size
-                )
+            if x.shape[1] % chunk_size != 0:
+                final_right_padding = chunk_size - (x.shape[1] % chunk_size)
             else:
                 final_right_padding = 0
 
             # compute the left context that can and should be added, for each
             # chunk. for the first few chunks, we will need to add extra padding
             applied_left_context = [
-                min(chunk_left_context, i * dynchunktrain_config.chunk_size,)
+                min(chunk_left_context, i * chunk_size)
                 for i in range(chunk_count)
             ]
 
@@ -190,10 +187,9 @@ def forward(
             out = [
                 x[
                     :,
-                    (
-                        i * dynchunktrain_config.chunk_size
-                        - applied_left_context[i]
-                    ) : ((i + 1) * dynchunktrain_config.chunk_size),
+                    (i * chunk_size - applied_left_context[i]) : (
+                        (i + 1) * chunk_size
+                    ),
                     ...,
                 ]
                 for i in range(chunk_count)

From 86850adeab7400b26b46520fb8d30640d167afe8 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 11:47:37 +0100
Subject: [PATCH 69/83] Add docstrings to forward_streaming methods

---
 .../lobes/models/transformer/Conformer.py     | 35 ++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 96dd64eca5..0e2edf23d1 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -455,6 +455,23 @@ def forward_streaming(
         context: ConformerEncoderLayerStreamingContext,
         pos_embs: torch.Tensor = None,
     ):
+        """Conformer layer streaming forward (typically for
+        DynamicChunkTraining-trained models), which is to be used at inference
+        time. Relies on a mutable context object as initialized by
+        `make_streaming_context` that should be used across chunks.  
+        Invoked by `ConformerEncoder.forward_streaming`.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Input tensor for this layer. Batching is supported as long as you
+            keep the context consistent.
+        context: ConformerEncoderStreamingContext
+            Mutable streaming context; the same object should be passed across
+            calls.
+        pos_embs: torch.Tensor, optional
+            Positional embeddings, if used."""
+
         orig_len = x.shape[-2]
         # ffn module
         x = x + 0.5 * self.ffn_module1(x)
@@ -649,10 +666,26 @@ def forward(
 
     def forward_streaming(
         self,
-        src,
+        src: torch.Tensor,
         context: ConformerEncoderStreamingContext,
         pos_embs: Optional[torch.Tensor] = None,
     ):
+        """Conformer streaming forward (typically for
+        DynamicChunkTraining-trained models), which is to be used at inference
+        time. Relies on a mutable context object as initialized by
+        `make_streaming_context` that should be used across chunks.
+
+        Arguments
+        ---------
+        src : torch.Tensor
+            Input tensor. Batching is supported as long as you keep the context
+            consistent.
+        context: ConformerEncoderStreamingContext
+            Mutable streaming context; the same object should be passed across
+            calls.
+        pos_embs: torch.Tensor, optional
+            Positional embeddings, if used."""
+
         if self.attention_type == "RelPosMHAXL":
             if pos_embs is None:
                 raise ValueError(

From 392ed08a01bf950f11bb1d6d3ccd82dd6df1e7ad Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 11:50:27 +0100
Subject: [PATCH 70/83] Add a reference on Dynamic Chunk Training

---
 speechbrain/utils/dynamic_chunk_training.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/speechbrain/utils/dynamic_chunk_training.py b/speechbrain/utils/dynamic_chunk_training.py
index be2c196961..47b0effd71 100644
--- a/speechbrain/utils/dynamic_chunk_training.py
+++ b/speechbrain/utils/dynamic_chunk_training.py
@@ -1,6 +1,10 @@
 """Configuration and utility classes for classes for Dynamic Chunk Training, as
 often used for the training of streaming-capable models in speech recognition.
 
+The definition of Dynamic Chunk Training is based on that of the following
+paper, though a lot of the literature refers to the same definition:
+https://arxiv.org/abs/2012.05481
+
 Authors
 * Sylvain de Langen 2023
 """

From 721f1472510c74081ccfee2ef886d070f11029e1 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 11:59:31 +0100
Subject: [PATCH 71/83] Rework conformer docstring docs

---
 .../lobes/models/transformer/Conformer.py     | 32 +++++++++++++++----
 1 file changed, 25 insertions(+), 7 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 0e2edf23d1..a395c5448e 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -31,7 +31,10 @@ class ConformerEncoderLayerStreamingContext:
     """Streaming metadata and state for a `ConformerEncoderLayer`.
 
     The multi-head attention and Dynamic Chunk Convolution require to save some
-    left context that gets inserted as left padding."""
+    left context that gets inserted as left padding.
+    
+    See :class:`.ConvolutionModule` documentation for further details.
+    """
 
     mha_left_context_size: int
     """For this layer, specifies how many frames of inputs should be saved.
@@ -144,15 +147,30 @@ def __init__(
 
     def forward(
         self,
-        x,
-        mask=None,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
         dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
     ):
-        """ Processes the input tensor x and returns the output an output tensor"""
+        """Applies the convolution to an input tensor `x`.
+
+        Arguments
+        ---------
+        x: torch.Tensor
+            Input tensor to the convolution module.
+        mask: torch.Tensor, optional
+            Mask to be applied over the output of the convolution using
+            `masked_fill_`, if specified.
+        dynchunktrain_config: DynChunkTrainConfig, optional
+            If specified, makes the module support Dynamic Chunk Convolution
+            (DCConv) as implemented by
+            `Dynamic Chunk Convolution for Unified Streaming and Non-Streaming Conformer ASR <https://www.amazon.science/publications/dynamic-chunk-convolution-for-unified-streaming-and-non-streaming-conformer-asr>`_.
+            This allows masking future frames while preserving better accuracy
+            than a fully causal convolution, at a small speed cost.
+            This should only be used for training (or, if you know what you're
+            doing, for masked evaluation at inference time), as the forward
+            streaming function should be used at inference time.
+            """
 
-        # ref: Dynamic chunk convolution for unified streaming and non-streaming
-        # conformer ASR
-        # https://www.amazon.science/publications/dynamic-chunk-convolution-for-unified-streaming-and-non-streaming-conformer-asr
         # split the input into chunks of size `chunk_size`, but for each chunk
         # provide a left context for left chunk dependencies to be possible.
 

From a459180b5ab4e44ef09badb126a51868bdd08e71 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 12:11:27 +0100
Subject: [PATCH 72/83] Update conformer author list, fix doc formatting for
 authors

---
 speechbrain/lobes/models/transformer/Conformer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index a395c5448e..87fcc43564 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -1,8 +1,10 @@
 """Conformer implementation.
 
 Authors
+-------
 * Jianyuan Zhong 2020
 * Samuele Cornell 2021
+* Sylvain de Langen 2023
 """
 
 from dataclasses import dataclass

From b2f6b5c0c1ededcab913b742777bf23e30568799 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 12:12:11 +0100
Subject: [PATCH 73/83] Fix trailing whitespace in conformer

---
 speechbrain/lobes/models/transformer/Conformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 87fcc43564..e4ff859efb 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -34,7 +34,7 @@ class ConformerEncoderLayerStreamingContext:
 
     The multi-head attention and Dynamic Chunk Convolution require to save some
     left context that gets inserted as left padding.
-    
+
     See :class:`.ConvolutionModule` documentation for further details.
     """
 
@@ -478,7 +478,7 @@ def forward_streaming(
         """Conformer layer streaming forward (typically for
         DynamicChunkTraining-trained models), which is to be used at inference
         time. Relies on a mutable context object as initialized by
-        `make_streaming_context` that should be used across chunks.  
+        `make_streaming_context` that should be used across chunks.
         Invoked by `ConformerEncoder.forward_streaming`.
 
         Arguments

From 86b8bba894f0d3b229fddbee6ad810ef5e28cc05 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 12:52:24 +0100
Subject: [PATCH 74/83] Improved comments in Conformer.forward

---
 .../lobes/models/transformer/Conformer.py     | 29 +++++++++++++------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index e4ff859efb..2d8e99076e 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -173,9 +173,6 @@ def forward(
             streaming function should be used at inference time.
             """
 
-        # split the input into chunks of size `chunk_size`, but for each chunk
-        # provide a left context for left chunk dependencies to be possible.
-
         if dynchunktrain_config is not None:
             # chances are chunking+causal is unintended; i don't know where it
             # may make sense, but if it does to you, feel free to implement it.
@@ -183,12 +180,23 @@ def forward(
                 not self.causal
             ), "Chunked convolution not supported with causal padding"
 
+            # in a causal convolution, which is not the case here, an output
+            # frame would never be able to depend on a input frame from any
+            # point in the future.
+
+            # but with the dynamic chunk convolution, we instead use a "normal"
+            # convolution but where, for any output frame, the future beyond the
+            # "current" chunk gets masked.
+            # see the paper linked in the documentation for details.
+
             chunk_size = dynchunktrain_config.chunk_size
             batch_size = x.shape[0]
             chunk_left_context = self.padding
 
             chunk_count = int(math.ceil(x.shape[1] / chunk_size))
 
+            # determine the amount of padding we need to insert at the right of
+            # the last chunk so that all chunks end up with the same size.
             if x.shape[1] % chunk_size != 0:
                 final_right_padding = chunk_size - (x.shape[1] % chunk_size)
             else:
@@ -202,7 +210,7 @@ def forward(
             ]
 
             # build views of chunks with left context (but no 0-padding yet)
-            # the left context effectively becomes "left padding", we do not
+            # the left context is treated as if it were left padding: we do not
             # want to keep any convolution results centered on the left context
             out = [
                 x[
@@ -215,6 +223,10 @@ def forward(
                 for i in range(chunk_count)
             ]
 
+            # TODO: it should be possible to insert some padding to stack all
+            # the tensors at this level. currently, this is rather inefficient
+            # as this as to be called on every individual chunk.
+
             out = [self.layer_norm(chk) for chk in out]
             out = [chk.transpose(1, 2) for chk in out]
             out = [self.bottleneck(chk) for chk in out]
@@ -260,11 +272,10 @@ def forward(
 
             # let's keep backwards compat by pointing at the weights from the
             # already declared Conv1d.
-
-            # we do not need to edit the bottleneck as it is pointwise (i.e.
-            # time step by time step), thus, it doesn't need padding along the
-            # time dimension
-
+            # in the prior steps, we manually applied:
+            # - left padding (known left context + zeroes if necessary)
+            # - right padding (zeroes)
+            # hence we're fully disabling conv1d's own padding.
             # -> [batch_size * num_chunks, out_channels, chunk_size + rpad]
             out = F.conv1d(
                 x,

From 9c63fe2a91a2cc04632dc522753894af9758854d Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 13:02:51 +0100
Subject: [PATCH 75/83] Added random dynchunktrain sampler example

---
 speechbrain/utils/dynamic_chunk_training.py | 27 ++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/speechbrain/utils/dynamic_chunk_training.py b/speechbrain/utils/dynamic_chunk_training.py
index 47b0effd71..5f0d8892d7 100644
--- a/speechbrain/utils/dynamic_chunk_training.py
+++ b/speechbrain/utils/dynamic_chunk_training.py
@@ -46,7 +46,32 @@ def is_infinite_left_context(self) -> bool:
 @dataclass
 class DynChunkTrainConfigRandomSampler:
     """Helper class to generate a DynChunkTrainConfig at runtime depending on the current
-    stage."""
+    stage.
+
+    Example
+    -------
+    >>> from speechbrain.core import Stage
+    >>> from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
+    >>> from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfigRandomSampler
+    >>> # for the purpose of this example, we test a scenario with a 100%
+    >>> # chance of the (24, None) scenario to occur
+    >>> sampler = DynChunkTrainConfigRandomSampler(
+    ...     chunkwise_prob=1.0,
+    ...     chunk_size_min=24,
+    ...     chunk_size_max=24,
+    ...     limited_left_context_prob=0.0,
+    ...     left_context_chunks_min=16,
+    ...     left_context_chunks_max=16,
+    ...     test_config=DynChunkTrainConfig(32, 16),
+    ...     valid_config=None
+    ... )
+    >>> one_train_config = sampler(Stage.TRAIN)
+    >>> one_train_config
+    DynChunkTrainConfig(chunk_size=24, left_context_size=None)
+    >>> one_train_config.is_infinite_left_context()
+    True
+    >>> sampler(Stage.TEST)
+    DynChunkTrainConfig(chunk_size=32, left_context_size=16)"""
 
     chunkwise_prob: float
     """When sampling (during `Stage.TRAIN`), the probability that a finite chunk

From eee3752417405df034f46c6d12961c4faf3ccbbf Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 13:04:08 +0100
Subject: [PATCH 76/83] More explicit names for mask functions in
 TransformerASR

---
 speechbrain/lobes/models/transformer/TransformerASR.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index dfc16c58ac..6518213d9d 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -36,7 +36,7 @@ class TransformerASRStreamingContext:
     """
 
 
-def make_asr_src_mask(
+def make_transformer_src_mask(
     src: torch.Tensor,
     causal: bool = False,
     dynchunktrain_config: Optional[DynChunkTrainConfig] = None,
@@ -113,7 +113,7 @@ def make_asr_src_mask(
     return None
 
 
-def make_asr_masks(
+def make_transformer_src_tgt_masks(
     src,
     tgt=None,
     wav_len=None,
@@ -146,7 +146,7 @@ def make_asr_masks(
         src_key_padding_mask = ~length_to_mask(abs_len).bool()
 
     # mask out the source
-    src_mask = make_asr_src_mask(
+    src_mask = make_transformer_src_mask(
         src, causal=causal, dynchunktrain_config=dynchunktrain_config
     )
 
@@ -327,7 +327,7 @@ def forward(self, src, tgt, wav_len=None, pad_idx=0):
             tgt_key_padding_mask,
             src_mask,
             tgt_mask,
-        ) = make_asr_masks(
+        ) = make_transformer_src_tgt_masks(
             src, tgt, wav_len, causal=self.causal, pad_idx=pad_idx
         )
 
@@ -438,7 +438,7 @@ def encode(
             bz, t, ch1, ch2 = src.shape
             src = src.reshape(bz, t, ch1 * ch2)
 
-        (src_key_padding_mask, _, src_mask, _,) = make_asr_masks(
+        (src_key_padding_mask, _, src_mask, _,) = make_transformer_src_tgt_masks(
             src,
             None,
             wav_len,

From 17d4f5f113177232c312e82c5f993d4c0ac9cfd3 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 13:22:22 +0100
Subject: [PATCH 77/83] Added docstring example on encode_streaming

---
 .../models/transformer/TransformerASR.py      | 39 +++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 6518213d9d..615ee4eb83 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -483,6 +483,45 @@ def encode_streaming(self, src, context: TransformerASRStreamingContext):
         Returns
         -------
         Encoder output for this chunk.
+
+        Example
+        -------
+        >>> import torch
+        >>> from speechbrain.lobes.models.transformer.TransformerASR import TransformerASR
+        >>> from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig
+        >>> net = TransformerASR(
+        ...     tgt_vocab=100,
+        ...     input_size=64,
+        ...     d_model=64,
+        ...     nhead=8,
+        ...     num_encoder_layers=1,
+        ...     num_decoder_layers=0,
+        ...     d_ffn=128,
+        ...     attention_type="RelPosMHAXL",
+        ...     positional_encoding=None,
+        ...     encoder_module="conformer",
+        ...     normalize_before=True,
+        ...     causal=False,
+        ... )
+        >>> ctx = net.make_streaming_context(
+        ...     DynChunkTrainConfig(16, 24),
+        ...     encoder_kwargs={"mha_left_context_size": 24},
+        ... )
+        >>> src1 = torch.rand([8, 16, 64])
+        >>> src2 = torch.rand([8, 16, 64])
+        >>> out1 = net.encode_streaming(src1, ctx)
+        >>> out1.shape
+        torch.Size([8, 16, 64])
+        >>> ctx.encoder_context.layers[0].mha_left_context.shape
+        torch.Size([8, 16, 64])
+        >>> out2 = net.encode_streaming(src2, ctx)
+        >>> out2.shape
+        torch.Size([8, 16, 64])
+        >>> ctx.encoder_context.layers[0].mha_left_context.shape
+        torch.Size([8, 24, 64])
+        >>> combined_out = torch.concat((out1, out2), dim=1)
+        >>> combined_out.shape
+        torch.Size([8, 32, 64])
         """
 
         if src.dim() == 4:

From ccc00f66b30732fba2cd2c3fe16306ebf210f01e Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 14:01:13 +0100
Subject: [PATCH 78/83] Pre-commit fix

---
 speechbrain/lobes/models/transformer/TransformerASR.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/TransformerASR.py b/speechbrain/lobes/models/transformer/TransformerASR.py
index 615ee4eb83..05a9d2e2be 100755
--- a/speechbrain/lobes/models/transformer/TransformerASR.py
+++ b/speechbrain/lobes/models/transformer/TransformerASR.py
@@ -438,7 +438,12 @@ def encode(
             bz, t, ch1, ch2 = src.shape
             src = src.reshape(bz, t, ch1 * ch2)
 
-        (src_key_padding_mask, _, src_mask, _,) = make_transformer_src_tgt_masks(
+        (
+            src_key_padding_mask,
+            _,
+            src_mask,
+            _,
+        ) = make_transformer_src_tgt_masks(
             src,
             None,
             wav_len,

From dbcdf9aaa40524ec46ccdb5fbdd08469386c4f53 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 14:51:16 +0100
Subject: [PATCH 79/83] Fix typo in conformer

---
 speechbrain/lobes/models/transformer/Conformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/speechbrain/lobes/models/transformer/Conformer.py b/speechbrain/lobes/models/transformer/Conformer.py
index 2d8e99076e..9220583021 100755
--- a/speechbrain/lobes/models/transformer/Conformer.py
+++ b/speechbrain/lobes/models/transformer/Conformer.py
@@ -278,7 +278,7 @@ def forward(
             # hence we're fully disabling conv1d's own padding.
             # -> [batch_size * num_chunks, out_channels, chunk_size + rpad]
             out = F.conv1d(
-                x,
+                out,
                 weight=self.conv.weight,
                 bias=self.conv.bias,
                 stride=self.conv.stride,

From eb67e1a04cff7071c2fe507513034c195d4aa899 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 15:54:28 +0100
Subject: [PATCH 80/83] Initial streaming integration test

---
 ...onformertransducer_streaming_experiment.py | 284 ++++++++++++++++++
 .../hyperparams.yaml                          | 268 +++++++++++++++++
 2 files changed, 552 insertions(+)
 create mode 100644 tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py
 create mode 100644 tests/integration/ASR_ConformerTransducer_streaming/hyperparams.yaml

diff --git a/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py b/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py
new file mode 100644
index 0000000000..2e832102c4
--- /dev/null
+++ b/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py
@@ -0,0 +1,284 @@
+#!/usr/bin/env/python3
+"""This minimal example trains a RNNT-based speech recognizer on a tiny dataset.
+The encoder is based on a Conformer model with the use of Dynamic Chunk Training
+ (with a Dynamic Chunk Convolution within the convolution modules) that predict
+phonemes. A greedy search is used on top of the output probabilities.
+Given the tiny dataset, the expected behavior is to overfit the training dataset
+(with a validation performance that stays high).
+"""
+import pathlib
+import speechbrain as sb
+from hyperpyyaml import load_hyperpyyaml
+import torch
+
+
+class ConformerTransducerBrain(sb.Brain):
+    def compute_forward(self, batch, stage):
+        """Forward computations from the waveform batches to the output probabilities."""
+        batch = batch.to(self.device)
+        wavs, wav_lens = batch.sig
+        tokens_with_bos, token_with_bos_lens = batch.phn_encoded_bos
+
+        # Add waveform augmentation if specified.
+        if stage == sb.Stage.TRAIN:
+            if hasattr(self.hparams, "wav_augment"):
+                wavs, wav_lens = self.hparams.wav_augment(wavs, wav_lens)
+                tokens_with_bos = self.hparams.wav_augment.replicate_labels(
+                    tokens_with_bos
+                )
+
+        feats = self.hparams.compute_features(wavs)
+
+        # Add feature augmentation if specified.
+        if stage == sb.Stage.TRAIN and hasattr(self.hparams, "fea_augment"):
+            feats, fea_lens = self.hparams.fea_augment(feats, wav_lens)
+            tokens_with_bos = self.hparams.fea_augment.replicate_labels(
+                tokens_with_bos
+            )
+
+        current_epoch = self.hparams.epoch_counter.current
+
+        # Old models may not have the streaming hparam, we don't break them in
+        # any other way so just check for its presence
+        if hasattr(self.hparams, "streaming") and self.hparams.streaming:
+            dynchunktrain_config = self.hparams.dynchunktrain_config_sampler(
+                stage
+            )
+        else:
+            dynchunktrain_config = None
+
+        feats = self.modules.normalize(feats, wav_lens, epoch=current_epoch)
+
+        src = self.modules.CNN(feats)
+        x = self.modules.enc(
+            src,
+            wav_lens,
+            pad_idx=self.hparams.pad_index,
+            dynchunktrain_config=dynchunktrain_config,
+        )
+        x = self.modules.proj_enc(x)
+
+        e_in = self.modules.emb(tokens_with_bos)
+        e_in = torch.nn.functional.dropout(
+            e_in,
+            self.hparams.dec_emb_dropout,
+            training=(stage == sb.Stage.TRAIN),
+        )
+        h, _ = self.modules.dec(e_in)
+        h = torch.nn.functional.dropout(
+            h, self.hparams.dec_dropout, training=(stage == sb.Stage.TRAIN)
+        )
+        h = self.modules.proj_dec(h)
+
+        # Joint network
+        # add labelseq_dim to the encoder tensor: [B,T,H_enc] => [B,T,1,H_enc]
+        # add timeseq_dim to the decoder tensor: [B,U,H_dec] => [B,1,U,H_dec]
+        joint = self.modules.Tjoint(x.unsqueeze(2), h.unsqueeze(1))
+
+        # Output layer for transducer log-probabilities
+        logits_transducer = self.modules.transducer_lin(joint)
+
+        # Compute outputs
+        if stage == sb.Stage.TRAIN:
+            p_ctc = None
+            p_ce = None
+
+            if self.hparams.ctc_weight > 0.0:
+                # Output layer for ctc log-probabilities
+                out_ctc = self.modules.proj_ctc(x)
+                p_ctc = self.hparams.log_softmax(out_ctc)
+
+            if self.hparams.ce_weight > 0.0:
+                # Output layer for ctc log-probabilities
+                p_ce = self.modules.dec_lin(h)
+                p_ce = self.hparams.log_softmax(p_ce)
+
+            return p_ctc, p_ce, logits_transducer, wav_lens
+
+        best_hyps, scores, _, _ = self.hparams.Greedysearcher(x)
+        return logits_transducer, wav_lens, best_hyps
+
+    def compute_objectives(self, predictions, batch, stage):
+        """Computes the loss (Transducer+(CTC+NLL)) given predictions and targets."""
+
+        ids = batch.id
+        tokens, token_lens = batch.phn_encoded
+        tokens_eos, token_eos_lens = batch.phn_encoded_eos
+
+        # Train returns 4 elements vs 3 for val and test
+        if len(predictions) == 4:
+            p_ctc, p_ce, logits_transducer, wav_lens = predictions
+        else:
+            logits_transducer, wav_lens, predicted_tokens = predictions
+
+        if stage == sb.Stage.TRAIN:
+            if hasattr(self.hparams, "wav_augment"):
+                tokens = self.hparams.wav_augment.replicate_labels(tokens)
+                token_lens = self.hparams.wav_augment.replicate_labels(
+                    token_lens
+                )
+                tokens_eos = self.hparams.wav_augment.replicate_labels(
+                    tokens_eos
+                )
+                token_eos_lens = self.hparams.wav_augment.replicate_labels(
+                    token_eos_lens
+                )
+            if hasattr(self.hparams, "fea_augment"):
+                tokens = self.hparams.fea_augment.replicate_labels(tokens)
+                token_lens = self.hparams.fea_augment.replicate_labels(
+                    token_lens
+                )
+                tokens_eos = self.hparams.fea_augment.replicate_labels(
+                    tokens_eos
+                )
+                token_eos_lens = self.hparams.fea_augment.replicate_labels(
+                    token_eos_lens
+                )
+
+        if stage == sb.Stage.TRAIN:
+            CTC_loss = 0.0
+            CE_loss = 0.0
+            if p_ctc is not None:
+                CTC_loss = self.hparams.ctc_cost(
+                    p_ctc, tokens, wav_lens, token_lens
+                )
+            if p_ce is not None:
+                CE_loss = self.hparams.ce_cost(
+                    p_ce, tokens_eos, length=token_eos_lens
+                )
+            loss_transducer = self.hparams.transducer_cost(
+                logits_transducer, tokens, wav_lens, token_lens
+            )
+            loss = (
+                self.hparams.ctc_weight * CTC_loss
+                + self.hparams.ce_weight * CE_loss
+                + (1 - (self.hparams.ctc_weight + self.hparams.ce_weight))
+                * loss_transducer
+            )
+        else:
+            loss = self.hparams.transducer_cost(
+                logits_transducer, tokens, wav_lens, token_lens
+            )
+
+        if stage != sb.Stage.TRAIN:
+            self.per_metrics.append(ids, predicted_tokens, tokens, target_len=token_lens)
+
+        return loss
+
+    def on_stage_start(self, stage, epoch=None):
+        "Gets called when a stage (either training, validation, test) starts."
+        if stage != sb.Stage.TRAIN:
+            self.per_metrics = self.hparams.per_stats()
+
+    def on_stage_end(self, stage, stage_loss, epoch=None):
+        """Gets called at the end of a stage."""
+        if stage == sb.Stage.TRAIN:
+            self.train_loss = stage_loss
+        if stage == sb.Stage.VALID and epoch is not None:
+            print("Epoch %d complete" % epoch)
+            print("Train loss: %.2f" % self.train_loss)
+        if stage != sb.Stage.TRAIN:
+            print(stage, "loss: %.2f" % stage_loss)
+            print(stage, "PER: %.2f" % self.per_metrics.summarize("error_rate"))
+
+
+def data_prep(data_folder, hparams):
+    "Creates the datasets and their data processing pipelines."
+
+    # 1. Declarations:
+    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
+        json_path=data_folder / "../annotation/ASR_train.json",
+        replacements={"data_root": data_folder},
+    )
+    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
+        json_path=data_folder / "../annotation/ASR_dev.json",
+        replacements={"data_root": data_folder},
+    )
+    datasets = [train_data, valid_data]
+    label_encoder = sb.dataio.encoder.CTCTextEncoder()
+    label_encoder.expect_len(hparams["num_labels"])
+
+    # 2. Define audio pipeline:
+    @sb.utils.data_pipeline.takes("wav")
+    @sb.utils.data_pipeline.provides("sig")
+    def audio_pipeline(wav):
+        sig = sb.dataio.dataio.read_audio(wav)
+        return sig
+
+    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)
+
+    # 3. Define text pipeline:
+    @sb.utils.data_pipeline.takes("phn")
+    @sb.utils.data_pipeline.provides(
+        "phn_list", "phn_encoded", "phn_encoded_bos", "phn_encoded_eos"
+    )
+    def text_pipeline(phn):
+        phn_list = phn.strip().split()
+        yield phn_list
+        phn_encoded = label_encoder.encode_sequence_torch(phn_list)
+        yield phn_encoded
+        phn_encoded_bos = label_encoder.prepend_bos_index(phn_encoded).long()
+        yield phn_encoded_bos
+        phn_encoded_eos = label_encoder.append_eos_index(phn_encoded).long()
+        yield phn_encoded_eos
+
+    sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline)
+
+    # 3. Fit encoder:
+    # NOTE: In this minimal example, also update from valid data
+    label_encoder.insert_blank(index=hparams["blank_index"])
+    label_encoder.insert_bos_eos(
+        bos_index=hparams["bos_index"], eos_label="<bos>"
+    )
+    label_encoder.update_from_didataset(train_data, output_key="phn_list")
+    label_encoder.update_from_didataset(valid_data, output_key="phn_list")
+
+    # 4. Set output:
+    sb.dataio.dataset.set_output_keys(
+        datasets, ["id", "sig", "phn_encoded", "phn_encoded_bos", "phn_encoded_eos"]
+    )
+    return train_data, valid_data, label_encoder
+
+
+def main(device="cpu"):
+    experiment_dir = pathlib.Path(__file__).resolve().parent
+    hparams_file = experiment_dir / "hyperparams.yaml"
+    data_folder = "../../samples/ASR"
+    data_folder = (experiment_dir / data_folder).resolve()
+
+    # Load model hyper parameters:
+    with open(hparams_file) as fin:
+        hparams = load_hyperpyyaml(fin)
+
+    # Dataset creation
+    train_data, valid_data, label_encoder = data_prep(data_folder, hparams)
+
+    # Trainer initialization
+    transducer_brain = ConformerTransducerBrain(
+        hparams["modules"],
+        hparams["opt_class"],
+        hparams,
+        run_opts={"device": device},
+    )
+
+    # Training/validation loop
+    transducer_brain.fit(
+        range(hparams["number_of_epochs"]),
+        train_data,
+        valid_data,
+        train_loader_kwargs=hparams["dataloader_options"],
+        valid_loader_kwargs=hparams["dataloader_options"],
+    )
+    # Evaluation is run separately (now just evaluating on valid data)
+    transducer_brain.evaluate(valid_data)
+
+    # Check that model overfits for integration test
+    assert transducer_brain.train_loss < 90.0
+
+
+if __name__ == "__main__":
+    main()
+
+
+def test_error(device):
+    main(device)
diff --git a/tests/integration/ASR_ConformerTransducer_streaming/hyperparams.yaml b/tests/integration/ASR_ConformerTransducer_streaming/hyperparams.yaml
new file mode 100644
index 0000000000..6d43121557
--- /dev/null
+++ b/tests/integration/ASR_ConformerTransducer_streaming/hyperparams.yaml
@@ -0,0 +1,268 @@
+# Seed needs to be set at top of yaml, before objects with parameters are made
+seed: 3407
+__set_seed: !!python/object/apply:torch.manual_seed [!ref <seed>]
+
+# Training parameters
+# To make Transformers converge, the global bath size should be large enough.
+# The global batch size is computed as batch_size * n_gpus * grad_accumulation_factor.
+# Empirically, we found that this value should be >= 128.
+# Please, set your parameters accordingly.
+number_of_epochs: 30
+lr: 1.0
+ctc_weight: 0.3 # Multitask with CTC for the encoder (0.0 = disabled)
+ce_weight: 0.0 # Multitask with CE for the decoder (0.0 = disabled)
+max_grad_norm: 5.0
+loss_reduction: 'batchmean'
+precision: fp32 # bf16, fp16 or fp32
+
+# Feature parameters
+sample_rate: 16000
+n_fft: 512
+n_mels: 80
+win_length: 32
+
+# Streaming & dynamic chunk training options
+# At least for the current architecture on LibriSpeech, we found out that
+# non-streaming accuracy is very similar between `streaming: True` and
+# `streaming: False`.
+streaming: True  # controls all Dynamic Chunk Training & chunk size & left context mechanisms
+
+# Configuration for Dynamic Chunk Training.
+# In this model, a chunk is roughly equivalent to 40ms of audio.
+dynchunktrain_config_sampler: !new:speechbrain.utils.dynamic_chunk_training.DynChunkTrainConfigRandomSampler # yamllint disable-line rule:line-length
+   chunkwise_prob: 0.6 # Probability during a batch to limit attention and sample a random chunk size in the following range
+   chunk_size_min: 2 # Minimum chunk size (if in a DynChunkTrain batch)
+   chunk_size_max: 8 # Maximum chunk size (if in a DynChunkTrain batch)
+   limited_left_context_prob: 0.75 # If in a DynChunkTrain batch, the probability during a batch to restrict left context to a random number of chunks
+   left_context_chunks_min: 1 # Minimum left context size (in # of chunks)
+   left_context_chunks_max: 8 # Maximum left context size (in # of chunks)
+   # If you specify a valid/test config, you can optionally have evaluation be
+   # done with a specific DynChunkTrain configuration.
+   # valid_config: !new:speechbrain.utils.dynamic_chunk_training.DynChunkTrainConfig
+   #    chunk_size: 24
+   #    left_context_size: 16
+   # test_config: ...
+
+dataloader_options:
+   batch_size: 1
+
+# Model parameters
+# Transformer
+d_model: 64
+joint_dim: 128
+nhead: 2
+num_encoder_layers: 1
+num_decoder_layers: 0
+d_ffn: 128
+transformer_dropout: 0.1
+activation: !name:torch.nn.GELU
+output_neurons: !ref <num_labels>
+dec_dim: 128
+dec_emb_dropout: 0.2
+dec_dropout: 0.1
+
+# Decoding parameters
+# Special tokens and labels
+blank_index: 0
+bos_index: 1
+pad_index: 1
+num_labels: 45
+beam_size: 10
+nbest: 1
+
+# If True uses torchaudio loss. Otherwise, the numba one
+use_torchaudio: True
+
+epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter
+   limit: !ref <number_of_epochs>
+
+normalize: !new:speechbrain.processing.features.InputNormalization
+   norm_type: global
+   update_until_epoch: 4
+
+compute_features: !new:speechbrain.lobes.features.Fbank
+   sample_rate: !ref <sample_rate>
+   n_fft: !ref <n_fft>
+   n_mels: !ref <n_mels>
+   win_length: !ref <win_length>
+
+# Speed perturbation
+speed_changes: [95, 100, 105]  # List of speed changes for time-stretching
+speed_perturb: !new:speechbrain.augment.time_domain.SpeedPerturb
+   orig_freq: !ref <sample_rate>
+   speeds: !ref <speed_changes>
+
+# Augmenter: Combines previously defined augmentations to perform data augmentation
+wav_augment: !new:speechbrain.augment.augmenter.Augmenter
+   parallel_augment: False
+   concat_original: False
+   repeat_augment: 1
+   shuffle_augmentations: False
+   min_augmentations: 1
+   max_augmentations: 1
+   augment_prob: 1.0
+   augmentations: [!ref <speed_perturb>]
+
+
+# Time Drop
+time_drop_length_low: 15  # Min length for temporal chunk to drop in spectrogram
+time_drop_length_high: 25  # Max length for temporal chunk to drop in spectrogram
+time_drop_count_low: 5  # Min number of chunks to drop in time in the spectrogram
+time_drop_count_high: 5  # Max number of chunks to drop in time in the spectrogram
+time_drop_replace: "zeros"  # Method of dropping chunks
+
+time_drop: !new:speechbrain.augment.freq_domain.SpectrogramDrop
+   drop_length_low: !ref <time_drop_length_low>
+   drop_length_high: !ref <time_drop_length_high>
+   drop_count_low: !ref <time_drop_count_low>
+   drop_count_high: !ref <time_drop_count_high>
+   replace: !ref <time_drop_replace>
+   dim: 1
+
+# Frequency Drop
+freq_drop_length_low: 25  # Min length for chunks to drop in frequency in the spectrogram
+freq_drop_length_high: 35  # Max length for chunks to drop in frequency in the spectrogram
+freq_drop_count_low: 2  # Min number of chunks to drop in frequency in the spectrogram
+freq_drop_count_high: 2  # Max number of chunks to drop in frequency in the spectrogram
+freq_drop_replace: "zeros"  # Method of dropping chunks
+
+freq_drop: !new:speechbrain.augment.freq_domain.SpectrogramDrop
+   drop_length_low: !ref <freq_drop_length_low>
+   drop_length_high: !ref <freq_drop_length_high>
+   drop_count_low: !ref <freq_drop_count_low>
+   drop_count_high: !ref <freq_drop_count_high>
+   replace: !ref <freq_drop_replace>
+   dim: 2
+
+# Time warp
+time_warp_window: 5  # Length of time warping window
+time_warp_mode: "bicubic"  # Time warping method
+
+time_warp: !new:speechbrain.augment.freq_domain.Warping
+   warp_window: !ref <time_warp_window>
+   warp_mode: !ref <time_warp_mode>
+   dim: 1
+
+fea_augment: !new:speechbrain.augment.augmenter.Augmenter
+   parallel_augment: False
+   concat_original: False
+   repeat_augment: 1
+   shuffle_augmentations: False
+   min_augmentations: 3
+   max_augmentations: 3
+   augment_prob: 1.0
+   augmentations: [
+      !ref <time_drop>,
+      !ref <freq_drop>,
+      !ref <time_warp>]
+
+CNN: !new:speechbrain.lobes.models.convolution.ConvolutionFrontEnd
+   input_shape: (8, 10, 80)
+   num_blocks: 2
+   num_layers_per_block: 1
+   out_channels: (64, 32)
+   kernel_sizes: (3, 3)
+   strides: (2, 2)
+   residuals: (False, False)
+
+Transformer: !new:speechbrain.lobes.models.transformer.TransformerASR.TransformerASR # yamllint disable-line rule:line-length
+   input_size: 640
+   tgt_vocab: !ref <output_neurons>
+   d_model: !ref <d_model>
+   nhead: !ref <nhead>
+   num_encoder_layers: !ref <num_encoder_layers>
+   num_decoder_layers: !ref <num_decoder_layers>
+   d_ffn: !ref <d_ffn>
+   dropout: !ref <transformer_dropout>
+   activation: !ref <activation>
+   encoder_module: conformer
+   attention_type: RelPosMHAXL
+   normalize_before: True
+   causal: False
+
+# We must call an encoder wrapper so the decoder isn't run (we don't have any)
+enc: !new:speechbrain.lobes.models.transformer.TransformerASR.EncoderWrapper
+   transformer: !ref <Transformer>
+
+# For MTL CTC over the encoder
+proj_ctc: !new:speechbrain.nnet.linear.Linear
+   input_size: !ref <joint_dim>
+   n_neurons: !ref <output_neurons>
+
+# Define some projection layers to make sure that enc and dec
+# output dim are the same before joining
+proj_enc: !new:speechbrain.nnet.linear.Linear
+   input_size: !ref <d_model>
+   n_neurons: !ref <joint_dim>
+   bias: False
+
+proj_dec: !new:speechbrain.nnet.linear.Linear
+   input_size: !ref <dec_dim>
+   n_neurons: !ref <joint_dim>
+   bias: False
+
+# Uncomment for MTL with CTC
+ctc_cost: !name:speechbrain.nnet.losses.ctc_loss
+   blank_index: !ref <blank_index>
+   reduction: !ref <loss_reduction>
+
+emb: !new:speechbrain.nnet.embedding.Embedding
+   num_embeddings: !ref <output_neurons>
+   consider_as_one_hot: True
+   blank_id: !ref <blank_index>
+
+dec: !new:speechbrain.nnet.RNN.LSTM
+   input_shape: [null, null, !ref <output_neurons> - 1]
+   hidden_size: !ref <dec_dim>
+   num_layers: 1
+   re_init: True
+
+# For MTL
+ce_cost: !name:speechbrain.nnet.losses.nll_loss
+   label_smoothing: 0.1
+
+Tjoint: !new:speechbrain.nnet.transducer.transducer_joint.Transducer_joint
+   joint: sum # joint [sum | concat]
+   nonlinearity: !ref <activation>
+
+transducer_lin: !new:speechbrain.nnet.linear.Linear
+   input_size: !ref <joint_dim>
+   n_neurons: !ref <output_neurons>
+   bias: False
+
+log_softmax: !new:speechbrain.nnet.activations.Softmax
+   apply_log: True
+
+transducer_cost: !name:speechbrain.nnet.losses.transducer_loss
+   blank_index: !ref <blank_index>
+   use_torchaudio: !ref <use_torchaudio>
+
+modules:
+   CNN: !ref <CNN>
+   enc: !ref <enc>
+   emb: !ref <emb>
+   dec: !ref <dec>
+   Tjoint: !ref <Tjoint>
+   transducer_lin: !ref <transducer_lin>
+   normalize: !ref <normalize>
+   proj_ctc: !ref <proj_ctc>
+   proj_dec: !ref <proj_dec>
+   proj_enc: !ref <proj_enc>
+
+model: !new:torch.nn.ModuleList
+   - [!ref <CNN>, !ref <enc>, !ref <emb>, !ref <dec>, !ref <proj_enc>, !ref <proj_dec>, !ref <proj_ctc>, !ref <transducer_lin>]
+
+Greedysearcher: !new:speechbrain.decoders.transducer.TransducerBeamSearcher
+   decode_network_lst: [!ref <emb>, !ref <dec>, !ref <proj_dec>]
+   tjoint: !ref <Tjoint>
+   classifier_network: [!ref <transducer_lin>]
+   blank_id: !ref <blank_index>
+   beam_size: 1
+   nbest: 1
+
+opt_class: !name:torch.optim.Adadelta
+    lr: !ref <lr>
+
+error_rate_computer: !name:speechbrain.utils.metric_stats.ErrorRateStats
+
+per_stats: !name:speechbrain.utils.metric_stats.ErrorRateStats

From f6213ecefb4ed7d39d472906a2077dd2952510d0 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 15:57:52 +0100
Subject: [PATCH 81/83] Precommit fix

---
 ...example_asr_conformertransducer_streaming_experiment.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py b/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py
index 2e832102c4..74f124706b 100644
--- a/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py
+++ b/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py
@@ -161,7 +161,9 @@ def compute_objectives(self, predictions, batch, stage):
             )
 
         if stage != sb.Stage.TRAIN:
-            self.per_metrics.append(ids, predicted_tokens, tokens, target_len=token_lens)
+            self.per_metrics.append(
+                ids, predicted_tokens, tokens, target_len=token_lens
+            )
 
         return loss
 
@@ -235,7 +237,8 @@ def text_pipeline(phn):
 
     # 4. Set output:
     sb.dataio.dataset.set_output_keys(
-        datasets, ["id", "sig", "phn_encoded", "phn_encoded_bos", "phn_encoded_eos"]
+        datasets,
+        ["id", "sig", "phn_encoded", "phn_encoded_bos", "phn_encoded_eos"],
     )
     return train_data, valid_data, label_encoder
 

From 56c69ff382728ab678dc2a3d8d52635499d01bd7 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 16:00:28 +0100
Subject: [PATCH 82/83] Fix indent in YAML

---
 .../ASR_ConformerTransducer_streaming/hyperparams.yaml          | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/integration/ASR_ConformerTransducer_streaming/hyperparams.yaml b/tests/integration/ASR_ConformerTransducer_streaming/hyperparams.yaml
index 6d43121557..a988432a34 100644
--- a/tests/integration/ASR_ConformerTransducer_streaming/hyperparams.yaml
+++ b/tests/integration/ASR_ConformerTransducer_streaming/hyperparams.yaml
@@ -261,7 +261,7 @@ Greedysearcher: !new:speechbrain.decoders.transducer.TransducerBeamSearcher
    nbest: 1
 
 opt_class: !name:torch.optim.Adadelta
-    lr: !ref <lr>
+   lr: !ref <lr>
 
 error_rate_computer: !name:speechbrain.utils.metric_stats.ErrorRateStats
 

From e70bb1de80f1b262f9ee69b2a5aef95184810a45 Mon Sep 17 00:00:00 2001
From: asu <sdelang@sdelang.fr>
Date: Mon, 18 Dec 2023 16:15:13 +0100
Subject: [PATCH 83/83] More consistent spelling in streaming integration test

---
 ...onformertransducer_streaming_experiment.py | 58 +++++++++----------
 1 file changed, 26 insertions(+), 32 deletions(-)

diff --git a/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py b/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py
index 74f124706b..893c511397 100644
--- a/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py
+++ b/tests/integration/ASR_ConformerTransducer_streaming/example_asr_conformertransducer_streaming_experiment.py
@@ -17,14 +17,14 @@ def compute_forward(self, batch, stage):
         """Forward computations from the waveform batches to the output probabilities."""
         batch = batch.to(self.device)
         wavs, wav_lens = batch.sig
-        tokens_with_bos, token_with_bos_lens = batch.phn_encoded_bos
+        phn_with_bos, phn_with_bos_lens = batch.phn_encoded_bos
 
         # Add waveform augmentation if specified.
         if stage == sb.Stage.TRAIN:
             if hasattr(self.hparams, "wav_augment"):
                 wavs, wav_lens = self.hparams.wav_augment(wavs, wav_lens)
-                tokens_with_bos = self.hparams.wav_augment.replicate_labels(
-                    tokens_with_bos
+                phn_with_bos = self.hparams.wav_augment.replicate_labels(
+                    phn_with_bos
                 )
 
         feats = self.hparams.compute_features(wavs)
@@ -32,8 +32,8 @@ def compute_forward(self, batch, stage):
         # Add feature augmentation if specified.
         if stage == sb.Stage.TRAIN and hasattr(self.hparams, "fea_augment"):
             feats, fea_lens = self.hparams.fea_augment(feats, wav_lens)
-            tokens_with_bos = self.hparams.fea_augment.replicate_labels(
-                tokens_with_bos
+            phn_with_bos = self.hparams.fea_augment.replicate_labels(
+                phn_with_bos
             )
 
         current_epoch = self.hparams.epoch_counter.current
@@ -58,7 +58,7 @@ def compute_forward(self, batch, stage):
         )
         x = self.modules.proj_enc(x)
 
-        e_in = self.modules.emb(tokens_with_bos)
+        e_in = self.modules.emb(phn_with_bos)
         e_in = torch.nn.functional.dropout(
             e_in,
             self.hparams.dec_emb_dropout,
@@ -102,52 +102,46 @@ def compute_objectives(self, predictions, batch, stage):
         """Computes the loss (Transducer+(CTC+NLL)) given predictions and targets."""
 
         ids = batch.id
-        tokens, token_lens = batch.phn_encoded
-        tokens_eos, token_eos_lens = batch.phn_encoded_eos
+        phn, phn_lens = batch.phn_encoded
+        phn_with_eos, phn_with_eos_lens = batch.phn_encoded_eos
 
         # Train returns 4 elements vs 3 for val and test
         if len(predictions) == 4:
             p_ctc, p_ce, logits_transducer, wav_lens = predictions
         else:
-            logits_transducer, wav_lens, predicted_tokens = predictions
+            logits_transducer, wav_lens, predicted_phn = predictions
 
         if stage == sb.Stage.TRAIN:
             if hasattr(self.hparams, "wav_augment"):
-                tokens = self.hparams.wav_augment.replicate_labels(tokens)
-                token_lens = self.hparams.wav_augment.replicate_labels(
-                    token_lens
+                phn = self.hparams.wav_augment.replicate_labels(phn)
+                phn_lens = self.hparams.wav_augment.replicate_labels(phn_lens)
+                phn_with_eos = self.hparams.wav_augment.replicate_labels(
+                    phn_with_eos
                 )
-                tokens_eos = self.hparams.wav_augment.replicate_labels(
-                    tokens_eos
-                )
-                token_eos_lens = self.hparams.wav_augment.replicate_labels(
-                    token_eos_lens
+                phn_with_eos_lens = self.hparams.wav_augment.replicate_labels(
+                    phn_with_eos_lens
                 )
             if hasattr(self.hparams, "fea_augment"):
-                tokens = self.hparams.fea_augment.replicate_labels(tokens)
-                token_lens = self.hparams.fea_augment.replicate_labels(
-                    token_lens
-                )
-                tokens_eos = self.hparams.fea_augment.replicate_labels(
-                    tokens_eos
+                phn = self.hparams.fea_augment.replicate_labels(phn)
+                phn_lens = self.hparams.fea_augment.replicate_labels(phn_lens)
+                phn_with_eos = self.hparams.fea_augment.replicate_labels(
+                    phn_with_eos
                 )
-                token_eos_lens = self.hparams.fea_augment.replicate_labels(
-                    token_eos_lens
+                phn_with_eos_lens = self.hparams.fea_augment.replicate_labels(
+                    phn_with_eos_lens
                 )
 
         if stage == sb.Stage.TRAIN:
             CTC_loss = 0.0
             CE_loss = 0.0
             if p_ctc is not None:
-                CTC_loss = self.hparams.ctc_cost(
-                    p_ctc, tokens, wav_lens, token_lens
-                )
+                CTC_loss = self.hparams.ctc_cost(p_ctc, phn, wav_lens, phn_lens)
             if p_ce is not None:
                 CE_loss = self.hparams.ce_cost(
-                    p_ce, tokens_eos, length=token_eos_lens
+                    p_ce, phn_with_eos, length=phn_with_eos_lens
                 )
             loss_transducer = self.hparams.transducer_cost(
-                logits_transducer, tokens, wav_lens, token_lens
+                logits_transducer, phn, wav_lens, phn_lens
             )
             loss = (
                 self.hparams.ctc_weight * CTC_loss
@@ -157,12 +151,12 @@ def compute_objectives(self, predictions, batch, stage):
             )
         else:
             loss = self.hparams.transducer_cost(
-                logits_transducer, tokens, wav_lens, token_lens
+                logits_transducer, phn, wav_lens, phn_lens
             )
 
         if stage != sb.Stage.TRAIN:
             self.per_metrics.append(
-                ids, predicted_tokens, tokens, target_len=token_lens
+                ids, predicted_phn, phn, target_len=phn_lens
             )
 
         return loss