Canary streaming

examples/asr/speech_to_text_eval.py

@@ -70,6 +70,7 @@
 import torch
 import transcribe_speech
 from omegaconf import MISSING, OmegaConf, open_dict
+from sacrebleu import corpus_bleu
 
 from nemo.collections.asr.metrics.wer import word_error_rate
 from nemo.collections.asr.parts.utils.transcribe_utils import (
@@ -140,19 +141,19 @@ def main(cfg: EvaluationConfig):
         transcription_cfg = cfg
 
     ground_truth_text = []
+    answers_text = []
     predicted_text = []
     invalid_manifest = False
     with open(transcription_cfg.output_filename, 'r') as f:
         for line in f:
             data = json.loads(line)
-
             if "pred_text" not in data:
                 invalid_manifest = True
                 break
-
             ground_truth_text.append(data[cfg.gt_text_attr_name])
-
             predicted_text.append(data["pred_text"])
+            if "answer" in data:
+                answers_text.append(data["answer"])
 
     pc = PunctuationCapitalization(cfg.text_processing.punctuation_marks)
     if cfg.text_processing.separate_punctuation:
@@ -213,6 +214,9 @@ def main(cfg: EvaluationConfig):
         logging.info(f'Got {metric_name} of {metric_value}. Tolerance was {cfg.tolerance}')
 
     logging.info(f"Dataset WER/CER {wer:.2%}/{cer:.2%}")
+    if answers_text:
+        bleu = corpus_bleu(predicted_text, [answers_text]).score
+        logging.info(f"Dataset BLEU {bleu:.2f}")
 
     if cfg.use_punct_er:
         dper_obj.print()

nemo/collections/asr/models/aed_multitask_models.py

@@ -910,6 +910,8 @@ def _transcribe_forward(
 
         log_probs, encoded_len, enc_states, enc_mask = self.forward(input_signal=audio, input_signal_length=audio_lens)
 
+        # import pdb; pdb.set_trace()
+
         if decoder_input_ids is None:
             # The dataloader provided only audio + audio_lens, so we
             # are constructing the prompt dynamically using TranscribeConfig.

nemo/collections/asr/modules/transformer/transformer_decoders.py

@@ -63,7 +63,11 @@ def __init__(
         )
         self.layer_norm_2 = nn.LayerNorm(hidden_size, eps=1e-5)
         self.second_sub_layer = MultiHeadAttention(
-            hidden_size, num_attention_heads, attn_score_dropout, attn_layer_dropout
+            hidden_size,
+            num_attention_heads,
+            attn_score_dropout,
+            attn_layer_dropout,
+            return_xatt_scores=True,
         )
         self.layer_norm_3 = nn.LayerNorm(hidden_size, eps=1e-5)
         self.third_sub_layer = PositionWiseFF(hidden_size, inner_size, ffn_dropout, hidden_act)
@@ -79,7 +83,7 @@ def forward_preln(self, decoder_query, decoder_mask, decoder_keys, encoder_state
         residual = decoder_query
         decoder_query = self.layer_norm_1(decoder_query)
         decoder_keys = self.layer_norm_1(decoder_keys)
-        self_attn_output = self.first_sub_layer(decoder_query, decoder_keys, decoder_keys, decoder_mask)
+        self_attn_output, _ = self.first_sub_layer(decoder_query, decoder_keys, decoder_keys, decoder_mask)
         self_attn_output += residual
 
         if self.is_adapter_available():
@@ -95,7 +99,9 @@ def forward_preln(self, decoder_query, decoder_mask, decoder_keys, encoder_state
 
         residual = self_attn_output
         self_attn_output = self.layer_norm_2(self_attn_output)
-        enc_dec_attn_output = self.second_sub_layer(self_attn_output, encoder_states, encoder_states, encoder_mask)
+        enc_dec_attn_output, extra_output = self.second_sub_layer(
+            self_attn_output, encoder_states, encoder_states, encoder_mask
+        )
         enc_dec_attn_output += residual
 
         residual = enc_dec_attn_output
@@ -112,14 +118,14 @@ def forward_preln(self, decoder_query, decoder_mask, decoder_keys, encoder_state
             pack_input = self.forward_enabled_adapters(pack_input)
             output_states = pack_input['x']
 
-        return output_states
+        return output_states, extra_output
 
     def forward_postln(self, decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask):
         """
         Post-LayerNorm block
         Order of operations: Self-Attn -> Residual -> LN -> Cross-Attn -> Residual -> LN -> FFN -> Residual -> LN
         """
-        self_attn_output = self.first_sub_layer(decoder_query, decoder_keys, decoder_keys, decoder_mask)
+        self_attn_output, _ = self.first_sub_layer(decoder_query, decoder_keys, decoder_keys, decoder_mask)
         self_attn_output += decoder_query
 
         if self.is_adapter_available():
@@ -135,7 +141,9 @@ def forward_postln(self, decoder_query, decoder_mask, decoder_keys, encoder_stat
 
         self_attn_output = self.layer_norm_1(self_attn_output)
 
-        enc_dec_attn_output = self.second_sub_layer(self_attn_output, encoder_states, encoder_states, encoder_mask)
+        enc_dec_attn_output, extra_output = self.second_sub_layer(
+            self_attn_output, encoder_states, encoder_states, encoder_mask
+        )
         enc_dec_attn_output += self_attn_output
         enc_dec_attn_output = self.layer_norm_2(enc_dec_attn_output)
 
@@ -151,7 +159,7 @@ def forward_postln(self, decoder_query, decoder_mask, decoder_keys, encoder_stat
             pack_ip = self.forward_enabled_adapters(pack_ip)
             output_states = pack_ip['x']
 
-        return self.layer_norm_3(output_states)
+        return self.layer_norm_3(output_states), xatt_scores
 
     def forward(self, decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask):
         if self.pre_ln:
@@ -251,9 +259,14 @@ def forward(
             else:
                 cached_mems_list = memory_states.unsqueeze(0)
 
+        xatt_scores_list = []
+
         for i, layer in enumerate(self.layers):
-            decoder_states = layer(decoder_states, decoder_attn_mask, memory_states, encoder_states, encoder_attn_mask)
+            decoder_states, extra_output = layer(
+                decoder_states, decoder_attn_mask, memory_states, encoder_states, encoder_attn_mask
+            )
             memory_states = self._get_memory_states(decoder_states, decoder_mems_list, i + 1)
+            xatt_scores_list.append(extra_output['xatt_scores'])
             if return_mems:
                 if return_mems_as_list:
                     cached_mems_list.append(memory_states)
@@ -270,9 +283,9 @@ def forward(
                     cached_mems_list = torch.cat((cached_mems_list, memory_states.unsqueeze(0)), dim=0)
 
         if return_mems:
-            return cached_mems_list
+            return cached_mems_list, xatt_scores_list
         else:
-            return memory_states
+            return memory_states, xatt_scores_list
 
     def input_example(self, max_batch=1, max_dim=256):
         """

nemo/collections/asr/modules/transformer/transformer_generators.py

@@ -23,6 +23,8 @@
 from nemo.collections.asr.parts.utils.asr_confidence_utils import ConfidenceMethodMixin
 from nemo.collections.common.parts import NEG_INF, mask_padded_tokens
 
+from nemo.utils import logging
+
 __all__ = [
     "GreedySequenceGenerator",
     "TopKSequenceGenerator",
@@ -128,6 +130,7 @@ def _one_step_forward(
         decoder_mems_list=None,
         pos=0,
         return_scores: bool = True,
+        return_xatt_scores: bool = False,
     ):
         """
         One step of autoregressive output generation.
@@ -148,21 +151,25 @@ def _one_step_forward(
         decoder_input_mask = mask_padded_tokens(decoder_input_ids, self.pad).float()
 
         if encoder_hidden_states is not None:
-            decoder_mems_list = self.decoder.forward(
+            decoder_mems_list, xatt_scores_list = self.decoder.forward(
                 decoder_hidden_states,
                 decoder_input_mask,
                 encoder_hidden_states,
                 encoder_input_mask,
                 decoder_mems_list,
                 return_mems=True,
             )
+            # import pdb; pdb.set_trace()
         else:
-            decoder_mems_list = self.decoder.forward(
+            decoder_mems_list, _ = self.decoder.forward(
                 decoder_hidden_states, decoder_input_mask, decoder_mems_list, return_mems=True
             )
         with self.classifier.with_log_softmax_enabled(return_scores) as clf:
             logits = clf.forward(hidden_states=decoder_mems_list[-1][:, -1:])
-        return logits, decoder_mems_list
+        if return_xatt_scores:
+            return logits, decoder_mems_list, xatt_scores_list
+        else:
+            return logits, decoder_mems_list
 
     def _prepare_for_search(self, decoder_input_ids=None, encoder_hidden_states=None):
         """
@@ -202,6 +209,7 @@ def _forward(
         is_sampling = self.temperature is not None and self.n_samples > 1
 
         tgt, batch_size, max_generation_length = self._prepare_for_search(decoder_input_ids, encoder_hidden_states)
+        tgt_len = tgt.size(-1)
         if is_sampling:
             tgt = torch.repeat_interleave(tgt, self.n_samples, dim=0)
             encoder_hidden_states = torch.repeat_interleave(encoder_hidden_states, self.n_samples, dim=0)
@@ -228,8 +236,15 @@ def _forward(
             if i == 0:
                 input_ids = tgt
             else:
+                i += tgt_len - 1
                 input_ids = tgt[:, -1:]
 
+            # logging.warning(f"Step {i}")
+            # logging.warning(f"tgt: {tgt}")
+            # logging.warning(f"input_ids: {input_ids}")
+            # if i == 14:
+            #     raise ValueError("Stop here")
+
             logits, decoder_mems_list = self._one_step_forward(
                 input_ids,
                 encoder_hidden_states,
@@ -413,6 +428,15 @@ def _forward(
         scores, prefixes = torch.topk(log_probs.permute(0, 2, 1), self.beam_size, dim=1)
         scores, prefixes = scores.view(-1, 1), prefixes.view(-1, 1)
 
+        # logging.warning(f"Step {0}")
+        # logging.warning(f"decoder_input_ids:   {decoder_input_ids}")
+        # logging.warning(f"tgt:   {tgt}")
+        # logging.warning(f"prefixes[:, -1:]: {prefixes[:, -1:]}")
+        # logging.warning("**********"*100)
+        # logging.warning(f"encoder_hidden_states.shape {encoder_hidden_states.shape}")
+        # logging.warning(f"encoder_hidden_states[0,35] {encoder_hidden_states[0,:20]    }")
+        # raise ValueError("Stop here")
+
         # repeat init target prefixes and cached memory states beam_size times
         prefixes = torch.cat((tgt.repeat(1, self.beam_size).view(-1, tgt.shape[1]), prefixes), dim=1)
         for j in range(len(decoder_mems_list)):
@@ -439,6 +463,10 @@ def _forward(
         tgt_len = tgt.size(-1)
         for i in range(tgt_len, max_generation_length + tgt_len):
 
+            # import pdb; pdb.set_trace()
+            # logging.warning(f"Step {i}")
+            # logging.warning(f"prefixes[:, -1:]: {prefixes[:, -1:]}")
+            # raise ValueError("Stop here")
             # mask all finished hypotheses to exclude them from beam
             pad_mask = pad_profile.repeat(1, self.beam_size)
 
@@ -448,6 +476,8 @@ def _forward(
             )
             scores_i, prefixes_i = torch.topk(log_probs[:, -1, :], self.beam_size, dim=-1)
 
+            # logging.warning(f"prefixes_i: {prefixes_i}")
+
             # for all prefixes ending with <eos> or <pad> replace generated
             # continuations with <pad>
             prefixes_i = self.pad * pad_mask + prefixes_i * (1 - pad_mask)

nemo/collections/asr/modules/transformer/transformer_modules.py

@@ -109,11 +109,28 @@ def forward(self, input_ids, token_type_ids=None, start_pos=0):
                 f"Input sequence is longer than maximum allowed sequence length for positional encoding. "
                 f"Got {seq_length} and {self.max_sequence_length}"
             )
+
+        # prepare position embedding for asynchronius decoding (canary streaming)
+        if torch.is_tensor(start_pos):
+            shift_pos = start_pos.unsqueeze(-1)
+            start_pos = 0
+        else:
+            shift_pos = None
+
         position_ids = torch.arange(
             start=start_pos, end=start_pos + seq_length, dtype=torch.long, device=input_ids.device
         )
         position_ids = position_ids.unsqueeze(0).repeat(input_ids.size(0), 1)
 
+        # import pdb; pdb.set_trace()
+
+        if torch.is_tensor(shift_pos):
+            # shift_pos is a tensor, so we need to add it to the position_ids
+            # and make sure that the resulting position_ids are within the
+            # range of the positional embedding
+            position_ids = position_ids + shift_pos
+            # position_ids = torch.clamp(position_ids, 0, self.max_sequence_length - 1)
+
         token_embeddings = self.token_embedding(input_ids)
         position_embeddings = self.position_embedding(position_ids)
         embeddings = token_embeddings + position_embeddings
@@ -140,7 +157,14 @@ class MultiHeadAttention(nn.Module):
             whole layer, but before layer normalization
     """
 
-    def __init__(self, hidden_size, num_attention_heads, attn_score_dropout=0.0, attn_layer_dropout=0.0):
+    def __init__(
+        self,
+        hidden_size,
+        num_attention_heads,
+        attn_score_dropout=0.0,
+        attn_layer_dropout=0.0,
+        return_xatt_scores=False,
+    ):
         super().__init__()
         if hidden_size % num_attention_heads != 0:
             raise ValueError(
@@ -160,6 +184,8 @@ def __init__(self, hidden_size, num_attention_heads, attn_score_dropout=0.0, att
         self.attn_dropout = nn.Dropout(attn_score_dropout)
         self.layer_dropout = nn.Dropout(attn_layer_dropout)
 
+        self.return_xatt_scores = return_xatt_scores
+
     def transpose_for_scores(self, x):
         new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attn_head_size)
         x = x.view(*new_x_shape)
@@ -179,6 +205,9 @@ def forward(self, queries, keys, values, attention_mask):
 
         # for numerical stability we pre-divide query and key by sqrt(sqrt(d))
         attention_scores = torch.matmul(query, key.transpose(-1, -2))
+
+        # import pdb; pdb.set_trace()
+
         if attention_mask is not None:
             attention_scores = attention_scores + attention_mask.to(attention_scores.dtype)
         attention_probs = torch.softmax(attention_scores, dim=-1)
@@ -193,7 +222,12 @@ def forward(self, queries, keys, values, attention_mask):
         # output projection
         output_states = self.out_projection(context)
         output_states = self.layer_dropout(output_states)
-        return output_states
+
+        extra_output = {}
+        if self.return_xatt_scores:
+            extra_output['xatt_scores'] = attention_probs
+
+        return output_states, extra_output
 
 
 class PositionWiseFF(nn.Module):