Merge pull request #99 from TensorSpeech/fix/beamsearch

Fix transducer beam search for longer sequence
TensorSpeech · Jan 6, 2021 · 8b9383c · 8b9383c
2 parents 4842078 + 3d5c7f3
commit 8b9383c
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/tensorflow_asr/models/transducer.py b/tensorflow_asr/models/transducer.py
@@ -591,7 +591,7 @@ def initialize_beam(dynamic=False):
             B = BeamHypothesis(
                 score=B.score.write(0, 0.0),
                 indices=B.indices.write(0, self.text_featurizer.blank),
-                prediction=B.prediction.write(0, tf.ones([total], dtype=tf.int32) * self.text_featurizer.blank),
+                prediction=B.prediction.write(0, tf.ones([total * 2], dtype=tf.int32) * self.text_featurizer.blank),
                 states=B.states.write(0, self.predict_net.get_initial_state())
             )
 
@@ -673,10 +673,7 @@ def false_fn():
 
                         b_score, b_indices, b_prediction, b_states, \
                             a_score, a_indices, a_prediction, a_states, A_i = tf.cond(
-                                tf.equal(pred, self.text_featurizer.blank),
-                                true_fn=true_fn,
-                                false_fn=false_fn
-                            )
+                                tf.equal(pred, self.text_featurizer.blank), true_fn=true_fn, false_fn=false_fn)
 
                         B = BeamHypothesis(score=b_score, indices=b_indices, prediction=b_prediction, states=b_states)
                         A = BeamHypothesis(score=a_score, indices=a_indices, prediction=a_prediction, states=a_states)