eos renorm + ilm eos handling

Author: mmz33

users/zeineldeen/models/lm/external_lm_decoder.py

@@ -376,6 +376,7 @@ def __init__(
         length_normalization=True,
         mask_layer_name=None,
         eos_cond_layer_name=None,
+        handle_eos_for_ilm=False,
         renorm_wo_eos=False,
     ):
         self.asr_decoder = copy.deepcopy(asr_decoder)
@@ -388,10 +389,41 @@ def __init__(
         self.length_normalization = length_normalization
         self.mask_layer_name = mask_layer_name
         self.eos_cond_layer_name = eos_cond_layer_name
+        self.handle_eos_for_ilm = handle_eos_for_ilm
         self.renorm_wo_eos = renorm_wo_eos
 
         self.network = None
 
+    def _handle_EOS(self, lm_net_out, lm_output_prob, prefix=""):
+        # so this means that eos prob is only used when the condition is true
+        lm_output_prob_wo_eos_ = lm_net_out.add_slice_layer(
+            f"{prefix}lm_output_prob_wo_eos_", lm_output_prob, axis="F", slice_start=1
+        )  # [B,V-1]
+        lm_output_prob_eos = lm_net_out.add_slice_layer(
+            f"{prefix}lm_output_prob_eos", lm_output_prob, axis="F", slice_start=0, slice_end=1
+        )  # [B,1]
+        if self.renorm_wo_eos:
+            lm_output_prob_eos_renorm = lm_net_out.add_activation_layer(
+                "lm_output_prob_eos_renorm", lm_output_prob_eos, activation="softmax"
+            )  # [B,V-1]
+        else:
+            lm_output_prob_eos_renorm = lm_output_prob_eos  # [B,V-1]
+        prob_1_const = lm_net_out.add_eval_layer(
+            f"{prefix}prob_1_const", lm_output_prob_eos, eval="tf.ones_like(source(0))"
+        )  # convert to ones
+        lm_output_prob_wo_eos = lm_net_out.add_generic_layer(
+            f"{prefix}lm_output_prob_wo_eos",
+            cls="concat",
+            source=[(prob_1_const, "F"), (lm_output_prob_eos_renorm, "F")],
+        )
+        lm_output_prob_cond = lm_net_out.add_switch_layer(
+            f"{prefix}lm_output_prob_cond",
+            condition=self.eos_cond_layer_name,
+            true_from=lm_output_prob,
+            false_from=lm_output_prob_wo_eos,
+        )
+        return lm_output_prob_cond
+
     def _create_external_lm_net(self) -> dict:
         lm_net_out = ReturnnNetwork()
 
@@ -438,27 +470,7 @@ def _create_external_lm_net(self) -> dict:
             )
 
         if self.eos_cond_layer_name:
-            # so this means that eos prob is only used when the condition is true
-            lm_output_prob_wo_eos_ = lm_net_out.add_slice_layer(
-                "lm_output_prob_wo_eos_", lm_output_prob, axis="F", slice_start=1
-            )  # [B,V-1]
-            lm_output_prob_eos = lm_net_out.add_slice_layer(
-                "lm_output_prob_eos", lm_output_prob, axis="F", slice_start=0, slice_end=1
-            )  # [B,1]
-            prob_1_const = lm_net_out.add_eval_layer(
-                "prob_1_const", lm_output_prob_eos, eval="tf.ones_like(source(0))"
-            )  # convert to ones
-            lm_output_prob_wo_eos = lm_net_out.add_generic_layer(
-                "lm_output_prob_wo_eos",
-                cls="concat",
-                source=[(prob_1_const, "F"), (lm_output_prob_wo_eos_, "F")],
-            )
-            lm_output_prob = lm_net_out.add_switch_layer(
-                "lm_output_prob_cond",
-                condition=self.eos_cond_layer_name,
-                true_from=lm_output_prob,
-                false_from=lm_output_prob_wo_eos,
-            )
+            lm_output_prob = self._handle_EOS(lm_net_out, lm_output_prob)
 
         fusion_str = "safe_log(source(0)) + {} * safe_log(source(1))".format(ext_lm_scale)  # shallow fusion
         fusion_source = [self.am_output_prob, lm_output_prob]
@@ -475,8 +487,15 @@ def _create_external_lm_net(self) -> dict:
                 raise ValueError("dec type: {} is not valid".format(self.dec_type))
 
             ilm_decoder.create_network()  # add ILM
+
+            if self.handle_eos_for_ilm:
+                assert self.eos_cond_layer_name
+                ilm_output_prob = self._handle_EOS(lm_net_out, ilm_decoder.output_prob_name, prefix="ilm_")
+            else:
+                ilm_output_prob = ilm_decoder.output_prob_name
+
             fusion_str += " - {} * safe_log(source(2))".format(self.prior_lm_opts["scale"])
-            fusion_source += [ilm_decoder.output_prob_name]
+            fusion_source += [ilm_output_prob]
 
         if self.ext_lm_opts.get("local_norm", False):
             fusion_str = f"{fusion_str} - tf.math.reduce_logsumexp({fusion_str}, axis=-1, keepdims=True)"