✍️ update transducer prediction and scripts

Author: nglehuy

examples/conformer/train_tpu_keras_subword_conformer.py

@@ -32,7 +32,7 @@
 
 parser.add_argument("--bs", type=int, default=None, help="Batch size per replica")
 
-parser.add_argument("--spx", type=int, default=50, help="Steps per execution for maximizing TPU performance")
+parser.add_argument("--spx", type=int, default=1, help="Steps per execution for maximizing TPU performance")
 
 parser.add_argument("--tpu_address", type=str, default=None, help="TPU address. Leave None on Colab")
 

examples/contextnet/train_tpu_keras_subword_contextnet.py

@@ -32,7 +32,7 @@
 
 parser.add_argument("--bs", type=int, default=None, help="Batch size per replica")
 
-parser.add_argument("--spx", type=int, default=50, help="Steps per execution for maximizing TPU performance")
+parser.add_argument("--spx", type=int, default=1, help="Steps per execution for maximizing TPU performance")
 
 parser.add_argument("--tpu_address", type=str, default=None, help="TPU address. Leave None on Colab")
 

tensorflow_asr/models/transducer.py

@@ -93,11 +93,10 @@ 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, maxlen=self.max_length)
+            mask = tf.sequence_mask(prediction_length, maxlen=tf.shape(outputs)[1])
             outputs = rnn["rnn"](outputs, training=training, mask=mask)
             outputs = outputs[0]
             if rnn["ln"] is not None: