⚡ add max length to predict network

Author: nglehuy

tensorflow_asr/datasets/asr_dataset.py

@@ -90,7 +90,7 @@ def update_lengths(self, max_lengths_prefix: str = None):
     # -------------------------------- ENTRIES -------------------------------------
 
     def read_entries(self):
-        if hasattr(self, 'entries') and len(self.entries) > 0: return
+        if hasattr(self, "entries") and len(self.entries) > 0: return
         self.entries = []
         for file_path in self.data_paths:
             print(f"Reading {file_path} ...")

tensorflow_asr/models/transducer.py

@@ -93,10 +93,11 @@ def call(self, inputs, training=False, **kwargs):
         # inputs has shape [B, U]
         # use tf.gather_nd instead of tf.gather for tflite conversion
         outputs, prediction_length = inputs
+        if not hasattr(self, "max_length"): self.max_length = shape_list(outputs)[-1]
         outputs = self.embed(outputs, training=training)
         outputs = self.do(outputs, training=training)
         for rnn in self.rnns:
-            mask = tf.sequence_mask(prediction_length)
+            mask = tf.sequence_mask(prediction_length, maxlen=self.max_length)
             outputs = rnn["rnn"](outputs, training=training, mask=mask)
             outputs = outputs[0]
             if rnn["ln"] is not None: